diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 810538fc1..4997ba916 100644 --- a/langgraph/pregel/__init__.py +++ b/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, diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 4687a5af6..7dfc738e5 100644 --- a/langgraph/pregel/io.py +++ b/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 diff --git a/tests/test_pregel.py b/tests/test_pregel.py index a681a1fbd..f811a167d 100644 --- a/tests/test_pregel.py +++ b/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" + } diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e91455bf0..692b51e73 100644 --- a/tests/test_pregel_async.py +++ b/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