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..cee4e16dc 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. @@ -1113,6 +1113,20 @@ 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() + if mode in stream_modes: + if subgraphs and isinstance(stream_mode, list): + yield (tuple(ns.split(NS_SEP)) if ns else (), mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif subgraphs: + yield (tuple(ns.split(NS_SEP)) if ns else (), 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( @@ -1155,6 +1169,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 @@ -1174,13 +1190,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) @@ -1235,13 +1245,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 @@ -1257,13 +1262,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( @@ -1287,6 +1286,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 +1303,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. @@ -1354,6 +1355,20 @@ 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() + if mode in stream_modes: + if subgraphs and isinstance(stream_mode, list): + yield (tuple(ns.split(NS_SEP)) if ns else (), mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif subgraphs: + yield (tuple(ns.split(NS_SEP)) if ns else (), 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( @@ -1404,6 +1419,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 @@ -1424,13 +1441,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) @@ -1487,13 +1499,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 @@ -1509,13 +1517,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 58441f831..87551b83f 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, str, Any]]) -> None: ... + def popleft(self) -> Tuple[str, 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, str, Any]]) -> None: + for stream, vv in zip(self.streams, tee(values, len(self.streams))): + stream.extend(vv) + + def popleft(self) -> Tuple[str, 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.""" @@ -178,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 ) @@ -218,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 @@ -266,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, @@ -310,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/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..269e2bbcc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -10897,7 +10897,13 @@ 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"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ] # get state without subgraphs outer_state = app.get_state(config) assert outer_state == StateSnapshot( @@ -11117,7 +11123,15 @@ def test_doubly_nested_graph_state( }, ) # resume - app.invoke(None, config, debug=True) + assert [c for c in app.stream(None, config, subgraphs=True)] == [ + ( + (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 ( app.get_state(config) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 9163bfe7f..b24848e6d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9338,7 +9338,15 @@ 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"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"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 +9566,15 @@ 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)] == [ + ( + (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 ( await app.aget_state(config)