From d30da72f6ec5787d43b751efffb70653b279cc8d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 27 Mar 2025 17:47:11 -0700 Subject: [PATCH] Reduce the number of channels created for each node by 50% - Used to be 2 channels per node, it is now one per node, which is the minimum - Now both hard edges, conditional edges, entrypoint and conditional entrypoint all use the same channel to trigger a node --- libs/langgraph/langgraph/graph/state.py | 48 +++++-------- libs/langgraph/tests/test_large_cases.py | 60 +++++++--------- .../langgraph/tests/test_large_cases_async.py | 69 +++++-------------- libs/langgraph/tests/test_pregel.py | 39 ++++------- libs/langgraph/tests/test_pregel_async.py | 39 ++++------- 5 files changed, 90 insertions(+), 165 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 304bfa834..750ed30df 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -799,11 +799,11 @@ class CompiledStateGraph(CompiledGraph): raise InvalidUpdateError(msg) # state updaters - write_entries: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry]] = [ + write_entries: tuple[Union[ChannelWriteEntry, ChannelWriteTupleEntry], ...] = ( ChannelWriteTupleEntry( mapper=_get_root if output_keys == ["__root__"] else _get_updates - ) - ] + ), + ) # add node and output channel if key == START: @@ -811,20 +811,14 @@ class CompiledStateGraph(CompiledGraph): tags=[TAG_HIDDEN], triggers=[START], channels=[START], - writers=[ - ChannelWrite( - write_entries, - tags=[TAG_HIDDEN], - ), - ], + writers=[ChannelWrite(write_entries, tags=[TAG_HIDDEN])], ) elif node is not None: input_schema = node.input if node else self.builder.schema 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) + branch_channel = CHANNEL_BRANCH_TO.format(key) self.channels[branch_channel] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( triggers=[branch_channel], @@ -836,13 +830,8 @@ class CompiledStateGraph(CompiledGraph): input_schema, self.builder.type_hints[input_schema], ), - writers=[ - # publish to this channel and state keys - ChannelWrite( - write_entries + [ChannelWriteEntry(key, key)], - tags=[TAG_HIDDEN], - ), - ], + # publish to state keys + writers=[ChannelWrite(write_entries, tags=[TAG_HIDDEN])], metadata=node.metadata, retry_policy=node.retry_policy, bound=node.runnable, @@ -852,21 +841,13 @@ class CompiledStateGraph(CompiledGraph): def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None: if isinstance(starts, str): - if starts == START: - channel_name = f"start:{end}" - # register channel - self.channels[channel_name] = EphemeralValue(Any) - # subscribe to channel - self.nodes[end].triggers.append(channel_name) - # publish to channel - self.nodes[START].writers.append( + # subscribe to start channel + if end != END: + self.nodes[starts].writers.append( ChannelWrite( - [ChannelWriteEntry(channel_name, START)], tags=[TAG_HIDDEN] + (ChannelWriteEntry(CHANNEL_BRANCH_TO.format(end), starts),) ) ) - elif end != END: - # subscribe to start channel - self.nodes[end].triggers.append(starts) elif end != END: channel_name = f"join:{'+'.join(starts)}:{end}" # register channel @@ -877,7 +858,7 @@ class CompiledStateGraph(CompiledGraph): for start in starts: self.nodes[start].writers.append( ChannelWrite( - [ChannelWriteEntry(channel_name, start)], tags=[TAG_HIDDEN] + (ChannelWriteEntry(channel_name, start),), tags=[TAG_HIDDEN] ) ) @@ -890,7 +871,7 @@ class CompiledStateGraph(CompiledGraph): if filtered := [p for p in packets if p != END]: writes = [ ( - ChannelWriteEntry(f"branch:to:{p}", start) + ChannelWriteEntry(CHANNEL_BRANCH_TO.format(p), start) if not isinstance(p, Send) else p ) @@ -1166,3 +1147,6 @@ def _get_schema( if k in channels and isinstance(channels[k], BaseChannel) }, ) + + +CHANNEL_BRANCH_TO = "branch:to:{}" diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 6616d4cbb..5d6ba5e37 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2483,7 +2483,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"), + "langgraph_triggers": ("branch:to:agent",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2542,7 +2542,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"), + "langgraph_triggers": ("branch:to:agent",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2585,7 +2585,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"), + "langgraph_triggers": ("branch:to:agent",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -4660,7 +4660,7 @@ def test_root_graph( content="result for query", name="search_api", tool_call_id="tool_call123", - id="00000000-0000-4000-8000-000000000037", + id="00000000-0000-4000-8000-000000000040", ) ] }, @@ -4683,7 +4683,7 @@ def test_root_graph( content="result for another", name="search_api", tool_call_id="tool_call456", - id="00000000-0000-4000-8000-000000000045", + id="00000000-0000-4000-8000-000000000049", ) ] }, @@ -5387,7 +5387,7 @@ def test_root_graph( "__root__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000078", + id="00000000-0000-4000-8000-000000000083", ), AIMessage( content="", @@ -5407,7 +5407,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), AIMessage( - content="an extra message", id="00000000-0000-4000-8000-000000000100" + content="an extra message", id="00000000-0000-4000-8000-000000000107" ), HumanMessage(content="what is weather in la"), ], @@ -5501,10 +5501,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "rewrite_query", "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ( - "branch:to:rewrite_query", - "start:rewrite_query", - ), + "triggers": ("branch:to:rewrite_query",), }, }, ), @@ -5535,10 +5532,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_one", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ( - "branch:to:retriever_one", - "rewrite_query", - ), + "triggers": ("branch:to:retriever_one",), }, }, ), @@ -5552,10 +5546,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_two", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ( - "branch:to:retriever_two", - "rewrite_query", - ), + "triggers": ("branch:to:retriever_two",), }, }, ), @@ -5617,7 +5608,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "query": "query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], }, - "triggers": ("branch:to:qa", "retriever_one", "retriever_two"), + "triggers": ("branch:to:qa",), }, }, ), @@ -6643,7 +6634,7 @@ def test_branch_then( "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ("branch:to:prepare", "start:prepare"), + "triggers": ("branch:to:prepare",), }, }, { @@ -7795,7 +7786,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -7990,7 +7981,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -8033,7 +8024,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -8082,7 +8073,7 @@ def test_nested_graph_state( "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -8516,7 +8507,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child_1", "langgraph_path": [PULL, AnyStr("child_1")], "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), parent_config=( @@ -8602,7 +8593,6 @@ def test_doubly_nested_graph_state( "langgraph_step": 1, "langgraph_triggers": [ "branch:to:child_1", - AnyStr("start:child_1"), ], }, created_at=AnyStr(), @@ -8650,7 +8640,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -8946,7 +8936,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -8985,7 +8975,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -9037,7 +9027,7 @@ def test_doubly_nested_graph_state( "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -9091,7 +9081,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), parent_config={ @@ -9146,7 +9136,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), parent_config={ @@ -9208,7 +9198,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), parent_config={ @@ -9270,7 +9260,7 @@ def test_doubly_nested_graph_state( AnyStr("child_1"), ], "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")], + "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), parent_config=None, diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 9fe7f2ac3..c5d4a0ffc 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -2300,11 +2300,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ( - "branch:to:agent", - "start:agent", - "tools", - ), + "langgraph_triggers": ("branch:to:agent",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2363,11 +2359,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ( - "branch:to:agent", - "start:agent", - "tools", - ), + "langgraph_triggers": ("branch:to:agent",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2410,11 +2402,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ( - "branch:to:agent", - "start:agent", - "tools", - ), + "langgraph_triggers": ("branch:to:agent",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -3895,10 +3883,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "rewrite_query", "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ( - "branch:to:rewrite_query", - "start:rewrite_query", - ), + "triggers": ("branch:to:rewrite_query",), }, }, ), @@ -3929,10 +3914,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_one", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ( - "branch:to:retriever_one", - "rewrite_query", - ), + "triggers": ("branch:to:retriever_one",), }, }, ), @@ -3946,10 +3928,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_two", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ( - "branch:to:retriever_two", - "rewrite_query", - ), + "triggers": ("branch:to:retriever_two",), }, }, ), @@ -4011,7 +3990,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "query": "query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], }, - "triggers": ("branch:to:qa", "retriever_one", "retriever_two"), + "triggers": ("branch:to:qa",), }, }, ), @@ -4486,10 +4465,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ( - "branch:to:prepare", - "start:prepare", - ), + "triggers": ("branch:to:prepare",), }, }, { @@ -4805,10 +4781,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ( - "branch:to:prepare", - "start:prepare", - ), + "triggers": ("branch:to:prepare",), }, }, { @@ -5363,7 +5336,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -5560,7 +5533,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -5603,7 +5576,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -5652,7 +5625,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "inner", "langgraph_path": [PULL, "inner"], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner", "outer_1"], + "langgraph_triggers": ["branch:to:inner"], "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), @@ -6090,7 +6063,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child_1", "langgraph_path": [PULL, AnyStr("child_1")], "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1", "start:child_1"], + "langgraph_triggers": [ + "branch:to:child_1", + ], }, created_at=AnyStr(), parent_config=( @@ -6178,7 +6153,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_step": 1, "langgraph_triggers": [ "branch:to:child_1", - "start:child_1", ], }, created_at=AnyStr(), @@ -6230,7 +6204,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_step": 2, "langgraph_triggers": [ "branch:to:child", - AnyStr("parent_1"), ], "langgraph_checkpoint_ns": AnyStr("child:"), }, @@ -6529,7 +6502,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6568,7 +6541,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6620,7 +6593,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_node": "child", "langgraph_path": [PULL, AnyStr("child")], "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")], + "langgraph_triggers": ["branch:to:child"], "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), @@ -6680,7 +6653,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_step": 1, "langgraph_triggers": [ "branch:to:child_1", - AnyStr("start:child_1"), ], }, created_at=AnyStr(), @@ -6738,7 +6710,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_step": 1, "langgraph_triggers": [ "branch:to:child_1", - AnyStr("start:child_1"), ], }, created_at=AnyStr(), @@ -6803,7 +6774,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_step": 1, "langgraph_triggers": [ "branch:to:child_1", - AnyStr("start:child_1"), ], }, created_at=AnyStr(), @@ -6868,7 +6838,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_step": 1, "langgraph_triggers": [ "branch:to:child_1", - AnyStr("start:child_1"), ], }, created_at=AnyStr(), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f86f78405..34549caf4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1187,15 +1187,13 @@ def test_pending_writes_resume( assert checkpoint is not None # should contain error from "two" expected_writes = [ - (AnyStr(), "one", "one"), (AnyStr(), "value", 2), (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), ] - assert len(checkpoint.pending_writes) == 3 + assert len(checkpoint.pending_writes) == 2 assert all(w in expected_writes for w in checkpoint.pending_writes) # both non-error pending writes come from same task non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] - assert non_error_writes[0][0] == non_error_writes[1][0] # error write is from the other task error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) assert error_write[0] != non_error_writes[0][0] @@ -1242,10 +1240,10 @@ def test_pending_writes_resume( "pending_sends": [], "versions_seen": { "one": { - "start:one": AnyVersion(), + "branch:to:one": AnyVersion(), }, "two": { - "start:two": AnyVersion(), + "branch:to:two": AnyVersion(), }, "__input__": {}, "__start__": { @@ -1254,19 +1252,17 @@ def test_pending_writes_resume( "__interrupt__": { "value": AnyVersion(), "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), }, }, "channel_versions": { - "one": AnyVersion(), - "two": AnyVersion(), "value": AnyVersion(), "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), }, - "channel_values": {"one": "one", "two": "two", "value": 6}, + "channel_values": {"value": 6}, }, metadata={ "parents": {}, @@ -1309,13 +1305,13 @@ def test_pending_writes_resume( "channel_versions": { "value": AnyVersion(), "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), }, "channel_values": { "value": 1, - "start:one": "__start__", - "start:two": "__start__", + "branch:to:one": "__start__", + "branch:to:two": "__start__", }, }, metadata={ @@ -1333,10 +1329,8 @@ def test_pending_writes_resume( } }, pending_writes=UnsortedSequence( - (AnyStr(), "one", "one"), (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), - (AnyStr(), "two", "two"), (AnyStr(), "value", 3), ), ) @@ -1369,8 +1363,8 @@ def test_pending_writes_resume( parent_config=None, pending_writes=UnsortedSequence( (AnyStr(), "value", 1), - (AnyStr(), "start:one", "__start__"), - (AnyStr(), "start:two", "__start__"), + (AnyStr(), "branch:to:one", "__start__"), + (AnyStr(), "branch:to:two", "__start__"), ), ) @@ -6876,10 +6870,7 @@ def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ( - "branch:to:call_model", - "start:call_model", - ), + "langgraph_triggers": ("branch:to:call_model",), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7f9ef7cf2..d50589a24 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2021,15 +2021,13 @@ async def test_pending_writes_resume( assert checkpoint is not None # should contain error from "two" expected_writes = [ - (AnyStr(), "one", "one"), (AnyStr(), "value", 2), (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), ] - assert len(checkpoint.pending_writes) == 3 + assert len(checkpoint.pending_writes) == 2 assert all(w in expected_writes for w in checkpoint.pending_writes) # both non-error pending writes come from same task non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] - assert non_error_writes[0][0] == non_error_writes[1][0] # error write is from the other task error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) assert error_write[0] != non_error_writes[0][0] @@ -2076,10 +2074,10 @@ async def test_pending_writes_resume( "pending_sends": [], "versions_seen": { "one": { - "start:one": AnyVersion(), + "branch:to:one": AnyVersion(), }, "two": { - "start:two": AnyVersion(), + "branch:to:two": AnyVersion(), }, "__input__": {}, "__start__": { @@ -2088,19 +2086,17 @@ async def test_pending_writes_resume( "__interrupt__": { "value": AnyVersion(), "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), }, }, "channel_versions": { - "one": AnyVersion(), - "two": AnyVersion(), "value": AnyVersion(), "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), }, - "channel_values": {"one": "one", "two": "two", "value": 6}, + "channel_values": {"value": 6}, }, metadata={ "parents": {}, @@ -2145,13 +2141,13 @@ async def test_pending_writes_resume( "channel_versions": { "value": AnyVersion(), "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), }, "channel_values": { "value": 1, - "start:one": "__start__", - "start:two": "__start__", + "branch:to:one": "__start__", + "branch:to:two": "__start__", }, }, metadata={ @@ -2171,10 +2167,8 @@ async def test_pending_writes_resume( } }, pending_writes=UnsortedSequence( - (AnyStr(), "one", "one"), (AnyStr(), "value", 2), (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), - (AnyStr(), "two", "two"), (AnyStr(), "value", 3), ), ) @@ -2207,8 +2201,8 @@ async def test_pending_writes_resume( parent_config=None, pending_writes=UnsortedSequence( (AnyStr(), "value", 1), - (AnyStr(), "start:one", "__start__"), - (AnyStr(), "start:two", "__start__"), + (AnyStr(), "branch:to:one", "__start__"), + (AnyStr(), "branch:to:two", "__start__"), ), ) @@ -7593,10 +7587,7 @@ async def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ( - "branch:to:call_model", - "start:call_model", - ), + "langgraph_triggers": ("branch:to:call_model",), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"),