From 80e442e13d4d3ff3b032d1cd13c2da5cc38705db Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 29 Aug 2024 12:00:02 -0700 Subject: [PATCH] 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)