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
This commit is contained in:
Nuno Campos
2025-03-27 17:47:11 -07:00
parent 7f079adfee
commit d30da72f6e
5 changed files with 90 additions and 165 deletions
+16 -32
View File
@@ -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:{}"
+25 -35
View File
@@ -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,
+19 -50
View File
@@ -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(),
+15 -24
View File
@@ -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:"),
+15 -24
View File
@@ -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:"),