Merge pull request #320 from langchain-ai/nc/17apr/fix-stream-events

Fix astream_events modifying output format unexpectedly

langgraph/pregel/__init__.py

@@ -512,6 +512,7 @@ class Pregel(
 
     def _defaults(
         self,
+        config: Optional[RunnableConfig] = None,
         *,
         stream_mode: Optional[StreamMode] = None,
         input_keys: Optional[Union[str, Sequence[str]]] = None,
@@ -542,9 +543,13 @@ class Pregel(
             validate_keys(input_keys, self.channels)
         interrupt_before_nodes = interrupt_before_nodes or self.interrupt_before_nodes
         interrupt_after_nodes = interrupt_after_nodes or self.interrupt_after_nodes
+        stream_mode = stream_mode if stream_mode is not None else self.stream_mode
+        if config is not None and config.get("configurable", {}).get(CONFIG_KEY_READ):
+            # if being called as a node in another graph, always use values mode
+            stream_mode = "values"
         return (
             debug,
-            stream_mode if stream_mode is not None else self.stream_mode,
+            stream_mode,
             input_keys,
             output_keys,
             interrupt_before_nodes,
@@ -567,7 +572,10 @@ class Pregel(
         config = ensure_config(config)
         callback_manager = get_callback_manager_for_config(config)
         run_manager = callback_manager.on_chain_start(
-            dumpd(self), input, name=config.get("run_name", self.get_name())
+            dumpd(self),
+            input,
+            name=config.get("run_name", self.get_name()),
+            run_id=config.get("run_id"),
         )
         try:
             if config["recursion_limit"] < 1:
@@ -581,6 +589,7 @@ class Pregel(
                 interrupt_before_nodes,
                 interrupt_after_nodes,
             ) = self._defaults(
+                config,
                 stream_mode=stream_mode,
                 input_keys=input_keys,
                 output_keys=output_keys,
@@ -763,7 +772,10 @@ class Pregel(
         config = ensure_config(config)
         callback_manager = get_async_callback_manager_for_config(config)
         run_manager = await callback_manager.on_chain_start(
-            dumpd(self), input, name=config.get("run_name", self.get_name())
+            dumpd(self),
+            input,
+            name=config.get("run_name", self.get_name()),
+            run_id=config.get("run_id"),
         )
         # if running from astream_log() run each proc with streaming
         do_stream = next(
@@ -786,6 +798,7 @@ class Pregel(
                 interrupt_before_nodes,
                 interrupt_after_nodes,
             ) = self._defaults(
+                config,
                 stream_mode=stream_mode,
                 input_keys=input_keys,
                 output_keys=output_keys,

langgraph/pregel/io.py

@@ -1,5 +1,7 @@
 from typing import Any, Iterator, Mapping, Optional, Sequence, Union
 
+from langchain_core.runnables.utils import AddableDict
+
 from langgraph.channels.base import BaseChannel, EmptyChannelError
 from langgraph.constants import TAG_HIDDEN
 from langgraph.pregel.log import logger
@@ -61,6 +63,14 @@ def map_input(
                 logger.warning(f"Input channel {k} not found in {input_channels}")
 
 
+class AddableValuesDict(AddableDict):
+    def __add__(self, other: dict[str, Any]) -> "AddableValuesDict":
+        return self | other
+
+    def __radd__(self, other: dict[str, Any]) -> "AddableValuesDict":
+        return other | self
+
+
 def map_output_values(
     output_channels: Union[str, Sequence[str]],
     pending_writes: Sequence[tuple[str, Any]],
@@ -72,7 +82,15 @@ def map_output_values(
             yield read_channel(channels, output_channels)
     else:
         if {c for c, _ in pending_writes if c in output_channels}:
-            yield read_channels(channels, output_channels)
+            yield AddableValuesDict(read_channels(channels, output_channels))
+
+
+class AddableUpdatesDict(AddableDict):
+    def __add__(self, other: dict[str, Any]) -> "AddableUpdatesDict":
+        return [self, other]
+
+    def __radd__(self, other: dict[str, Any]) -> "AddableUpdatesDict":
+        raise TypeError("AddableUpdatesDict does not support right-side addition")
 
 
 def map_output_updates(
@@ -84,17 +102,21 @@ def map_output_updates(
         t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
     ]
     if isinstance(output_channels, str):
-        if updated := {
-            node: value
-            for node, _, _, writes, _ in output_tasks
-            for chan, value in writes
-            if chan == output_channels
-        }:
+        if updated := AddableUpdatesDict(
+            {
+                node: value
+                for node, _, _, writes, _ in output_tasks
+                for chan, value in writes
+                if chan == output_channels
+            }
+        ):
             yield updated
     else:
-        if updated := {
-            node: {chan: value for chan, value in writes if chan in output_channels}
-            for node, _, _, writes, _ in output_tasks
-            if any(chan in output_channels for chan, _ in writes)
-        }:
+        if updated := AddableUpdatesDict(
+            {
+                node: {chan: value for chan, value in writes if chan in output_channels}
+                for node, _, _, writes, _ in output_tasks
+                if any(chan in output_channels for chan, _ in writes)
+            }
+        ):
             yield updated

tests/test_pregel.py

@@ -3622,3 +3622,55 @@ def test_simple_multi_edge() -> None:
       +---------+  """
     )
     assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
+
+
+def test_nested_graph() -> None:
+    class State(TypedDict):
+        my_key: str
+
+    def up(state: State):
+        return {"my_key": state["my_key"] + " there"}
+
+    inner = StateGraph(State)
+    inner.add_node("up", up)
+    inner.set_entry_point("up")
+    inner.set_finish_point("up")
+
+    def side(state: State):
+        return {"my_key": state["my_key"] + " and back again"}
+
+    graph = StateGraph(State)
+    graph.add_node("inner", inner.compile())
+    graph.add_node("side", side)
+    graph.set_entry_point("inner")
+    graph.add_edge("inner", "side")
+    graph.set_finish_point("side")
+
+    app = graph.compile()
+
+    assert app.get_graph().draw_ascii() == (
+        """+-----------+  
+| __start__ |  
++-----------+  
+      *        
+      *        
+      *        
+  +-------+    
+  | inner |    
+  +-------+    
+      *        
+      *        
+      *        
+  +------+     
+  | side |     
+  +------+     
+      *        
+      *        
+      *        
+ +---------+   
+ | __end__ |   
+ +---------+   """
+    )
+    assert app.invoke({"my_key": "my value"}) == {
+        "my_key": "my value there and back again"
+    }

tests/test_pregel_async.py

@@ -12,6 +12,7 @@ from typing import (
     TypedDict,
     Union,
 )
+from uuid import UUID
 
 import pytest
 from langchain_core.runnables import RunnableLambda, RunnablePassthrough
@@ -3310,3 +3311,95 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N
         },
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
+
+
+async def test_nested_graph() -> None:
+    class State(TypedDict):
+        my_key: str
+
+    async def up(state: State):
+        return {"my_key": state["my_key"] + " there"}
+
+    inner = StateGraph(State)
+    inner.add_node("up", up)
+    inner.set_entry_point("up")
+    inner.set_finish_point("up")
+
+    async def side(state: State):
+        return {"my_key": state["my_key"] + " and back again"}
+
+    graph = StateGraph(State)
+    graph.add_node("inner", inner.compile())
+    graph.add_node("side", side)
+    graph.set_entry_point("inner")
+    graph.add_edge("inner", "side")
+    graph.set_finish_point("side")
+
+    app = graph.compile()
+
+    assert app.get_graph().draw_ascii() == (
+        """+-----------+  
+| __start__ |  
++-----------+  
+      *        
+      *        
+      *        
+  +-------+    
+  | inner |    
+  +-------+    
+      *        
+      *        
+      *        
+  +------+     
+  | side |     
+  +------+     
+      *        
+      *        
+      *        
+ +---------+   
+ | __end__ |   
+ +---------+   """
+    )
+    assert await app.ainvoke({"my_key": "my value"}) == {
+        "my_key": "my value there and back again"
+    }
+    assert [chunk async for chunk in app.astream({"my_key": "my value"})] == [
+        {"inner": {"my_key": "my value there"}},
+        {"side": {"my_key": "my value there and back again"}},
+    ]
+    assert [
+        chunk
+        async for chunk in app.astream({"my_key": "my value"}, stream_mode="values")
+    ] == [
+        {"my_key": "my value"},
+        {"my_key": "my value there"},
+        {"my_key": "my value there and back again"},
+    ]
+    times_called = 0
+    async for event in app.astream_events(
+        {"my_key": "my value"},
+        version="v1",
+        config={"run_id": UUID(int=0)},
+        stream_mode="values",
+    ):
+        if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
+            times_called += 1
+            assert event["data"] == {
+                "output": {"my_key": "my value there and back again"}
+            }
+    assert times_called == 1
+    times_called = 0
+    async for event in app.astream_events(
+        {"my_key": "my value"},
+        version="v1",
+        config={"run_id": UUID(int=0)},
+    ):
+        if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
+            times_called += 1
+            assert event["data"] == {
+                "output": [
+                    {"inner": {"my_key": "my value there"}},
+                    {"side": {"my_key": "my value there and back again"}},
+                ]
+            }
+    assert times_called == 1