diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 0f064f518..e9963f5c8 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -74,7 +74,7 @@ def map_command( ) -> Iterator[tuple[str, str, Any]]: """Map input chunk to a sequence of pending writes in the form (channel, value).""" if cmd.graph == Command.PARENT: - raise InvalidUpdateError("There is not parent graph") + raise InvalidUpdateError("There is no parent graph") if cmd.goto: if isinstance(cmd.goto, (tuple, list)): sends = cmd.goto diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index b360a7440..1c9ea1c8f 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -2,6 +2,7 @@ import asyncio import concurrent.futures from collections import defaultdict, deque from contextlib import AsyncExitStack, ExitStack +from dataclasses import replace from inspect import signature from types import TracebackType from typing import ( @@ -66,6 +67,7 @@ from langgraph.errors import ( EmptyInputError, GraphDelegate, GraphInterrupt, + ParentCommand, ) from langgraph.managed.base import ( ManagedValueMapping, @@ -736,6 +738,16 @@ class PregelLoop(LoopProtocol): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # add current state to parent command + if isinstance(exc_value, ParentCommand): + cmd = exc_value.args[0] + state = ( + [(self.output_keys, read_channels(self.channels, self.output_keys))] + if isinstance(self.output_keys, str) + else list(read_channels(self.channels, self.output_keys).items()) + ) + exc_value.args = (replace(cmd, update=[*state, *cmd._update_as_tuples()]),) + # suppress interrupt suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested if suppress: # emit one last "values" event, with pending writes applied diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index eda5c7b51..b08bc88dc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4828,6 +4828,13 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) "source": "loop", "writes": { "alice": { + "messages": [ + _AnyIdHumanMessage( + content="get user name", + additional_kwargs={}, + response_metadata={}, + ), + ], "user_name": "Meow", } }, @@ -6167,6 +6174,88 @@ def test_multiple_subgraphs_checkpointer( ] +def test_merging_updates_command_parent(): + # simple reducer + def append_unique(left, right): + combined = list(left) + for item in right: + if item in combined: + continue + else: + combined.append(item) + return combined + + class State(TypedDict): + foo: str + bar: Annotated[list[str], append_unique] + + # Define subgraph + def subgraph_node_1(state: State): + return Command( + goto="subgraph_node_2", + update={ + "foo": "foo", + "bar": ["subgraph_node_1"], + }, + ) + + def subgraph_node_2(state: State): + return Command( + goto="node_3", + update={"bar": ["subgraph_node_2"]}, + graph=Command.PARENT, + ) + + subgraph_builder = StateGraph(State) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.add_node(subgraph_node_2) + subgraph_builder.add_edge(START, "subgraph_node_1") + + # Define main graph + def node_1(state: State): + return Command( + goto="node_2", + update={"bar": ["node_1"]}, + ) + + def node_3(state: State, store): + return Command( + update={"bar": ["node_3"]}, + ) + + main_builder = StateGraph(State) + main_builder.add_node("node_1", node_1) + main_builder.add_node("node_2", subgraph_builder.compile()) + main_builder.add_node("node_3", node_3) + main_builder.add_edge(START, "node_1") + main_builder.add_edge("node_2", "node_3") + main_graph = main_builder.compile() + + assert main_graph.invoke({"foo": ""}) == { + "foo": "foo", + "bar": ["node_1", "subgraph_node_1", "subgraph_node_2", "node_3"], + } + + assert list( + main_graph.stream({"foo": ""}, stream_mode="updates", subgraphs=True) + ) == [ + ((), {"node_1": {"bar": ["node_1"]}}), + ( + (AnyStr("node_2:"),), + {"subgraph_node_1": {"foo": "foo", "bar": ["subgraph_node_1"]}}, + ), + ( + (), + { + "node_2": [ + {"foo": "foo"}, + {"bar": ["node_1", "subgraph_node_1"]}, + {"bar": ["subgraph_node_2"]}, + ] + }, + ), + ((), {"node_3": {"bar": ["node_3"]}}), + ] def test_entrypoint_output_schema_with_return_and_save() -> None: """Test output schema inference with entrypoint.final.""" diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 38627d12d..5db7f6e57 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6148,6 +6148,13 @@ async def test_parent_command(checkpointer_name: str) -> None: "source": "loop", "writes": { "alice": { + "messages": [ + _AnyIdHumanMessage( + content="get user name", + additional_kwargs={}, + response_metadata={}, + ), + ], "user_name": "Meow", } },