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

@@ -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"

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

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."""

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():

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)

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)