Merge pull request #1534 from langchain-ai/nc/29aug/stream-subgraphs

Stream output from 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.
@@ -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(

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
 

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,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)

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)