From 80e442e13d4d3ff3b032d1cd13c2da5cc38705db Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 29 Aug 2024 12:00:02 -0700 Subject: [PATCH 1/3] Stream output from subgraphs - enabled by new argument stream(subgraphs=True) - the same stream_mode requested for parent graph is applied to all subgraphs --- libs/langgraph/langgraph/constants.py | 4 ++- libs/langgraph/langgraph/pregel/__init__.py | 28 +++++++++++-------- libs/langgraph/langgraph/pregel/loop.py | 31 ++++++++++++++++++++- libs/langgraph/tests/any_str.py | 1 - libs/langgraph/tests/test_pregel.py | 12 ++++++-- libs/langgraph/tests/test_pregel_async.py | 14 ++++++++-- 6 files changed, 72 insertions(+), 18 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index eeb320ed3..0a748a032 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -5,10 +5,12 @@ INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" -CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map" +CONFIG_KEY_STREAM = "__pregel_stream" CONFIG_KEY_STORE = "__pregel_store" CONFIG_KEY_RESUMING = "__pregel_resuming" CONFIG_KEY_TASK_ID = "__pregel_task_id" +# this one part of public API so more readable +CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map" INTERRUPT = "__interrupt__" ERROR = "__error__" TASKS = "__pregel_tasks" diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 235606929..3d676a7bc 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -63,6 +63,7 @@ from langgraph.constants import ( CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, + CONFIG_KEY_STREAM, ERROR, INTERRUPT, NS_END, @@ -990,18 +991,17 @@ class Pregel( def _defaults( self, - config: Optional[RunnableConfig] = None, + config: RunnableConfig, *, - stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, - output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before: Optional[Union[All, Sequence[str]]] = None, - interrupt_after: Optional[Union[All, Sequence[str]]] = None, - debug: Optional[bool] = None, + stream_mode: Optional[Union[StreamMode, list[StreamMode]]], + output_keys: Optional[Union[str, Sequence[str]]], + interrupt_before: Optional[Union[All, Sequence[str]]], + interrupt_after: Optional[Union[All, Sequence[str]]], + debug: Optional[bool], ) -> tuple[ bool, Sequence[StreamMode], Union[str, Sequence[str]], - Union[str, Sequence[str]], Optional[Sequence[str]], Optional[Sequence[str]], Optional[BaseCheckpointSaver], @@ -1016,12 +1016,10 @@ class Pregel( stream_mode = stream_mode if stream_mode is not None else self.stream_mode if not isinstance(stream_mode, list): stream_mode = [stream_mode] - if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None: + if CONFIG_KEY_READ in config.get("configurable", {}): # if being called as a node in another graph, always use values mode stream_mode = ["values"] - if config is not None and config.get("configurable", {}).get( - CONFIG_KEY_CHECKPOINTER - ): + if CONFIG_KEY_CHECKPOINTER in config.get("configurable", {}): checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][ CONFIG_KEY_CHECKPOINTER ] @@ -1046,6 +1044,7 @@ class Pregel( interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: Optional[bool] = None, + subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: """Stream graph steps for a single input. @@ -1062,6 +1061,7 @@ class Pregel( interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. debug: Whether to print debug information during execution, defaults to False. + subgraphs: Whether to stream subgraphs, defaults to False. Yields: The output of each step in the graph. The output shape depends on the stream_mode. @@ -1155,6 +1155,8 @@ class Pregel( output_keys=output_keys, stream_keys=self.stream_channels_asis, ) as loop: + if subgraphs: + loop.config["configurable"][CONFIG_KEY_STREAM] = loop.stream # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates # channel updates from step N are only visible in step N+1 @@ -1287,6 +1289,7 @@ class Pregel( interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: Optional[bool] = None, + subgraphs: bool = False, ) -> AsyncIterator[Union[dict[str, Any], Any]]: """Stream graph steps for a single input. @@ -1303,6 +1306,7 @@ class Pregel( interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. debug: Whether to print debug information during execution, defaults to False. + subgraphs: Whether to stream subgraphs, defaults to False. Yields: The output of each step in the graph. The output shape depends on the stream_mode. @@ -1404,6 +1408,8 @@ class Pregel( output_keys=output_keys, stream_keys=self.stream_channels_asis, ) as loop: + if subgraphs: + loop.config["configurable"][CONFIG_KEY_STREAM] = loop.stream aioloop = asyncio.get_event_loop() # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 58441f831..7ff95cc38 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -2,16 +2,19 @@ import asyncio import concurrent.futures from collections import deque from contextlib import AsyncExitStack, ExitStack +from itertools import tee from types import TracebackType from typing import ( Any, AsyncContextManager, Callable, ContextManager, + Iterable, List, Literal, Mapping, Optional, + Protocol, Sequence, Tuple, Type, @@ -39,6 +42,7 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, + CONFIG_KEY_STREAM, ERROR, INPUT, INTERRUPT, @@ -86,6 +90,27 @@ INPUT_RESUMING = object() EMPTY_SEQ = () +class StreamProtocol(Protocol): + def extend(self, values: Iterable[Tuple[str, Any]]) -> None: ... + def popleft(self) -> Tuple[str, Any]: ... + def __bool__(self) -> bool: ... + + +class DuplexStream(StreamProtocol): + def __init__(self, *streams: StreamProtocol) -> None: + self.streams = streams + + def extend(self, values: Iterable[Tuple[str, Any]]) -> None: + for stream, vv in zip(self.streams, tee(values, len(self.streams))): + stream.extend(vv) + + def popleft(self) -> Tuple[str, Any]: + return self.streams[0].popleft() + + def __bool__(self) -> bool: + return bool(self.streams[0]) + + class PregelLoop: input: Optional[Any] config: RunnableConfig @@ -127,7 +152,7 @@ class PregelLoop: "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" ] tasks: Sequence[PregelExecutableTask] - stream: deque[Tuple[str, Any]] + stream: StreamProtocol output: Union[None, dict[str, Any], Any] = None # public @@ -154,6 +179,10 @@ class PregelLoop: self.output_keys = output_keys self.stream_keys = stream_keys self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) + if CONFIG_KEY_STREAM in config["configurable"]: + self.stream = DuplexStream( + self.stream, config["configurable"][CONFIG_KEY_STREAM] + ) def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: """Put writes for a task, to be read by the next tick.""" diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index 28a67ddf6..32383df34 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -18,7 +18,6 @@ class AnyDict(dict): super().__init__(*args, **kwargs) def __eq__(self, other: object) -> bool: - print("did we get here") if not isinstance(other, dict) or len(self) != len(other): return False for k, v in self.items(): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 9b6a1967d..645761881 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -10897,7 +10897,10 @@ def test_doubly_nested_graph_state( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - app.invoke({"my_key": "my value"}, config, debug=True) + assert [c for c in app.stream({"my_key": "my value"}, config, subgraphs=True)] == [ + {"parent_1": {"my_key": "hi my value"}}, + {"grandchild_1": {"my_key": "hi my value here"}}, + ] # get state without subgraphs outer_state = app.get_state(config) assert outer_state == StateSnapshot( @@ -11117,7 +11120,12 @@ def test_doubly_nested_graph_state( }, ) # resume - app.invoke(None, config, debug=True) + assert [c for c in app.stream(None, config, subgraphs=True)] == [ + {"grandchild_2": {"my_key": "hi my value here and there"}}, + {"child_1": {"my_key": "hi my value here and there"}}, + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] # get state with and without subgraphs assert ( app.get_state(config) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 9163bfe7f..f59003ac4 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9338,7 +9338,12 @@ async def test_doubly_nested_graph_state( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - await app.ainvoke({"my_key": "my value"}, config, debug=True) + assert [ + c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) + ] == [ + {"parent_1": {"my_key": "hi my value"}}, + {"grandchild_1": {"my_key": "hi my value here"}}, + ] # get state without subgraphs outer_state = await app.aget_state(config) assert outer_state == StateSnapshot( @@ -9558,7 +9563,12 @@ async def test_doubly_nested_graph_state( }, ) # resume - await app.ainvoke(None, config, debug=True) + assert [c async for c in app.astream(None, config, subgraphs=True)] == [ + {"grandchild_2": {"my_key": "hi my value here and there"}}, + {"child_1": {"my_key": "hi my value here and there"}}, + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] # get state with and without subgraphs assert ( await app.aget_state(config) From b479b88c7e2912bebbbc78484f48af1e03a48218 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 29 Aug 2024 13:20:28 -0700 Subject: [PATCH 2/3] Add name of subgraph to streaming output --- libs/langgraph/langgraph/pregel/__init__.py | 91 +++++++++++---------- libs/langgraph/langgraph/pregel/loop.py | 21 +++-- libs/langgraph/tests/test_pregel.py | 12 +-- libs/langgraph/tests/test_pregel_async.py | 12 +-- 4 files changed, 73 insertions(+), 63 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3d676a7bc..515a33aa4 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1113,6 +1113,25 @@ class Pregel( {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} ``` """ + + def output() -> Iterator: + while loop.stream: + ns, mode, payload = loop.stream.popleft() + ns = ( + NS_SEP.join(p.split(NS_END)[0] for p in ns.split(NS_SEP)) + if ns + else "" + ) + if mode in stream_modes: + if subgraphs and isinstance(stream_mode, list): + yield (ns, mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif subgraphs: + yield (ns, payload) + else: + yield payload + config = ensure_config(merge_configs(self.config, config)) callback_manager = get_callback_manager_for_config(config) run_manager = callback_manager.on_chain_start( @@ -1176,13 +1195,7 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + yield from output() # debug flag if debug: print_step_tasks(loop.step, loop.tasks) @@ -1237,13 +1250,8 @@ class Pregel( # remove references to loop vars del fut, task # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + yield from output() + # maybe stop other tasks if _should_stop_others(done): break @@ -1259,13 +1267,7 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + yield from output() # handle exit if loop.status == "out_of_steps": raise GraphRecursionError( @@ -1358,6 +1360,25 @@ class Pregel( {'type': 'task_result', 'timestamp': '2024-06-23T...+00:00', 'step': 2, 'payload': {'id': '...', 'name': 'b', 'result': [('alist', ['there'])]}} ``` """ + + def output() -> Iterator: + while loop.stream: + ns, mode, payload = loop.stream.popleft() + ns = ( + NS_SEP.join(p.split(NS_END)[0] for p in ns.split(NS_SEP)) + if ns + else "" + ) + if mode in stream_modes: + if subgraphs and isinstance(stream_mode, list): + yield (ns, mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif subgraphs: + yield (ns, payload) + else: + yield payload + config = ensure_config(merge_configs(self.config, config)) callback_manager = get_async_callback_manager_for_config(config) run_manager = await callback_manager.on_chain_start( @@ -1430,13 +1451,8 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + for o in output(): + yield o # debug flag if debug: print_step_tasks(loop.step, loop.tasks) @@ -1493,13 +1509,9 @@ class Pregel( # remove references to loop vars del fut, task # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + for o in output(): + yield o + # maybe stop other tasks if _should_stop_others(done): break @@ -1515,13 +1527,8 @@ class Pregel( self.stream_channels_list, ) # emit output - while loop.stream: - mode, payload = loop.stream.popleft() - if mode in stream_modes: - if isinstance(stream_mode, list): - yield (mode, payload) - else: - yield payload + for o in output(): + yield o # handle exit if loop.status == "out_of_steps": raise GraphRecursionError( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ff95cc38..87551b83f 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -91,8 +91,8 @@ EMPTY_SEQ = () class StreamProtocol(Protocol): - def extend(self, values: Iterable[Tuple[str, Any]]) -> None: ... - def popleft(self) -> Tuple[str, Any]: ... + def extend(self, values: Iterable[Tuple[str, str, Any]]) -> None: ... + def popleft(self) -> Tuple[str, str, Any]: ... def __bool__(self) -> bool: ... @@ -100,11 +100,11 @@ class DuplexStream(StreamProtocol): def __init__(self, *streams: StreamProtocol) -> None: self.streams = streams - def extend(self, values: Iterable[Tuple[str, Any]]) -> None: + def extend(self, values: Iterable[Tuple[str, str, Any]]) -> None: for stream, vv in zip(self.streams, tee(values, len(self.streams))): stream.extend(vv) - def popleft(self) -> Tuple[str, Any]: + def popleft(self) -> Tuple[str, str, Any]: return self.streams[0].popleft() def __bool__(self) -> bool: @@ -207,11 +207,11 @@ class PregelLoop: ) if task := next((t for t in self.tasks if t.id == task_id), None): self.stream.extend( - ("updates", v) + (self.config["configurable"].get("checkpoint_ns", ""), "updates", v) for v in map_output_updates(self.output_keys, [(task, writes)]) ) self.stream.extend( - ("debug", v) + (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_task_results( self.step, [(task, writes)], self.stream_keys ) @@ -247,7 +247,7 @@ class PregelLoop: self._update_mv(key, values) # produce values output self.stream.extend( - ("values", v) + (self.config["configurable"].get("checkpoint_ns", ""), "values", v) for v in map_output_values(self.output_keys, writes, self.channels) ) # clear pending writes @@ -295,7 +295,7 @@ class PregelLoop: # produce debug output if self._checkpointer_put_after_previous is not None: self.stream.extend( - ("debug", v) + (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_checkpoint( self.step - 1, # printing checkpoint for previous step self.checkpoint_config, @@ -339,7 +339,10 @@ class PregelLoop: return False # produce debug output - self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks)) + self.stream.extend( + (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) + for v in map_debug_tasks(self.step, self.tasks) + ) return True diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 645761881..a6a338a75 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -10898,8 +10898,8 @@ def test_doubly_nested_graph_state( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert [c for c in app.stream({"my_key": "my value"}, config, subgraphs=True)] == [ - {"parent_1": {"my_key": "hi my value"}}, - {"grandchild_1": {"my_key": "hi my value here"}}, + ("", {"parent_1": {"my_key": "hi my value"}}), + ("child|child_1", {"grandchild_1": {"my_key": "hi my value here"}}), ] # get state without subgraphs outer_state = app.get_state(config) @@ -11121,10 +11121,10 @@ def test_doubly_nested_graph_state( ) # resume assert [c for c in app.stream(None, config, subgraphs=True)] == [ - {"grandchild_2": {"my_key": "hi my value here and there"}}, - {"child_1": {"my_key": "hi my value here and there"}}, - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ("child|child_1", {"grandchild_2": {"my_key": "hi my value here and there"}}), + ("child", {"child_1": {"my_key": "hi my value here and there"}}), + ("", {"child": {"my_key": "hi my value here and there"}}), + ("", {"parent_2": {"my_key": "hi my value here and there and back again"}}), ] # get state with and without subgraphs assert ( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f59003ac4..830d69238 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9341,8 +9341,8 @@ async def test_doubly_nested_graph_state( assert [ c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) ] == [ - {"parent_1": {"my_key": "hi my value"}}, - {"grandchild_1": {"my_key": "hi my value here"}}, + ("", {"parent_1": {"my_key": "hi my value"}}), + ("child|child_1", {"grandchild_1": {"my_key": "hi my value here"}}), ] # get state without subgraphs outer_state = await app.aget_state(config) @@ -9564,10 +9564,10 @@ async def test_doubly_nested_graph_state( ) # resume assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - {"grandchild_2": {"my_key": "hi my value here and there"}}, - {"child_1": {"my_key": "hi my value here and there"}}, - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ("child|child_1", {"grandchild_2": {"my_key": "hi my value here and there"}}), + ("child", {"child_1": {"my_key": "hi my value here and there"}}), + ("", {"child": {"my_key": "hi my value here and there"}}), + ("", {"parent_2": {"my_key": "hi my value here and there and back again"}}), ] # get state with and without subgraphs assert ( From ddf67d9233c296945aed2f611ec81c8458b54296 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 29 Aug 2024 13:32:21 -0700 Subject: [PATCH 3/3] Add ns to subgraph stream events --- libs/langgraph/langgraph/pregel/__init__.py | 18 ++++-------------- libs/langgraph/tests/test_pregel.py | 18 ++++++++++++------ libs/langgraph/tests/test_pregel_async.py | 18 ++++++++++++------ 3 files changed, 28 insertions(+), 26 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 515a33aa4..cee4e16dc 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1117,18 +1117,13 @@ class Pregel( def output() -> Iterator: while loop.stream: ns, mode, payload = loop.stream.popleft() - ns = ( - NS_SEP.join(p.split(NS_END)[0] for p in ns.split(NS_SEP)) - if ns - else "" - ) if mode in stream_modes: if subgraphs and isinstance(stream_mode, list): - yield (ns, mode, payload) + yield (tuple(ns.split(NS_SEP)) if ns else (), mode, payload) elif isinstance(stream_mode, list): yield (mode, payload) elif subgraphs: - yield (ns, payload) + yield (tuple(ns.split(NS_SEP)) if ns else (), payload) else: yield payload @@ -1364,18 +1359,13 @@ class Pregel( def output() -> Iterator: while loop.stream: ns, mode, payload = loop.stream.popleft() - ns = ( - NS_SEP.join(p.split(NS_END)[0] for p in ns.split(NS_SEP)) - if ns - else "" - ) if mode in stream_modes: if subgraphs and isinstance(stream_mode, list): - yield (ns, mode, payload) + yield (tuple(ns.split(NS_SEP)) if ns else (), mode, payload) elif isinstance(stream_mode, list): yield (mode, payload) elif subgraphs: - yield (ns, payload) + yield (tuple(ns.split(NS_SEP)) if ns else (), payload) else: yield payload diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a6a338a75..269e2bbcc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -10898,8 +10898,11 @@ def test_doubly_nested_graph_state( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert [c for c in app.stream({"my_key": "my value"}, config, subgraphs=True)] == [ - ("", {"parent_1": {"my_key": "hi my value"}}), - ("child|child_1", {"grandchild_1": {"my_key": "hi my value here"}}), + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), ] # get state without subgraphs outer_state = app.get_state(config) @@ -11121,10 +11124,13 @@ def test_doubly_nested_graph_state( ) # resume assert [c for c in app.stream(None, config, subgraphs=True)] == [ - ("child|child_1", {"grandchild_2": {"my_key": "hi my value here and there"}}), - ("child", {"child_1": {"my_key": "hi my value here and there"}}), - ("", {"child": {"my_key": "hi my value here and there"}}), - ("", {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ((AnyStr("child:"),), {"child_1": {"my_key": "hi my value here and there"}}), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), ] # get state with and without subgraphs assert ( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 830d69238..b24848e6d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9341,8 +9341,11 @@ async def test_doubly_nested_graph_state( assert [ c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) ] == [ - ("", {"parent_1": {"my_key": "hi my value"}}), - ("child|child_1", {"grandchild_1": {"my_key": "hi my value here"}}), + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), ] # get state without subgraphs outer_state = await app.aget_state(config) @@ -9564,10 +9567,13 @@ async def test_doubly_nested_graph_state( ) # resume assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - ("child|child_1", {"grandchild_2": {"my_key": "hi my value here and there"}}), - ("child", {"child_1": {"my_key": "hi my value here and there"}}), - ("", {"child": {"my_key": "hi my value here and there"}}), - ("", {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ((AnyStr("child:"),), {"child_1": {"my_key": "hi my value here and there"}}), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), ] # get state with and without subgraphs assert (