From 0baaf5cd964b9a13ce54e357a1897d6773a49725 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 30 Mar 2024 19:41:50 -0700 Subject: [PATCH] Remove private api to mutate values during call to stream() --- langgraph/pregel/__init__.py | 36 ++++++------------------------- tests/test_pregel.py | 42 ------------------------------------ tests/test_pregel_async.py | 36 ------------------------------- 3 files changed, 6 insertions(+), 108 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index f170dc746..2845cda90 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -43,7 +43,6 @@ from langchain_core.runnables.utils import ( ) from langchain_core.tracers.log_stream import LogStreamCallbackHandler -from langgraph.channels.any_value import AnyValue from langgraph.channels.base import ( AsyncChannelsManager, BaseChannel, @@ -52,8 +51,6 @@ from langgraph.channels.base import ( InvalidUpdateError, create_checkpoint, ) -from langgraph.channels.ephemeral_value import EphemeralValue -from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -522,12 +519,15 @@ class Pregel( # if no more tasks, we're done if not next_tasks: - break + if step == 0: + raise ValueError("No tasks to run in graph.") + else: + break elif step == config["recursion_limit"]: raise GraphRecursionError( f"Recursion limit of {config['recursion_limit']} reached" - "without hitting a stop condition. You can increase the limit" - "by setting the `recursion_limit` config key." + "without hitting a stop condition. You can increase the " + "limit by setting the `recursion_limit` config key." ) if debug: @@ -582,10 +582,6 @@ class Pregel( # yield current value and checkpoint view if step_output := map_output(output_keys, pending_writes, channels): yield step_output - # we can detect updates when output is multiple channels (ie. dict) - if not isinstance(output_keys, str): - # if view was updated, apply writes to channels - _apply_writes_from_view(checkpoint, channels, step_output) # with previous step's checkpoint if _should_interrupt( @@ -773,10 +769,6 @@ class Pregel( # yield current value and checkpoint view if step_output := map_output(output_keys, pending_writes, channels): yield step_output - # we can detect updates when output is multiple channels (ie. dict) - if not isinstance(output_keys, str): - # if view was updated, apply writes to channels - _apply_writes_from_view(checkpoint, channels, step_output) # with previous step's checkpoint if _should_interrupt( @@ -1074,22 +1066,6 @@ def _apply_writes( channels[chan].update([]) -def _apply_writes_from_view( - checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], values: dict[str, Any] -) -> None: - # Apply writes to channels - for chan, value in values.items(): - if value == _read_channel(channels, chan): - continue - - assert isinstance(channels[chan], (LastValue, EphemeralValue, AnyValue)), ( - f"Can't modify channel {chan} of type " - f"{channels[chan].__class__.__name__}" - ) - checkpoint["channel_versions"][chan] += 1 - channels[chan].update([values[chan]]) - - def _prepare_next_tasks( checkpoint: Checkpoint, processes: Mapping[str, ChannelInvoke], diff --git a/tests/test_pregel.py b/tests/test_pregel.py index ff8814e00..238a83670 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -179,29 +179,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: with pytest.raises(GraphRecursionError): app.invoke(2, {"recursion_limit": 1}) - for step, values in enumerate(app.stream(2), start=1): - if step == 1: - assert values == { - "inbox": 3, - } - elif step == 2: - assert values == { - "output": 4, - } - - for step, values in enumerate(app.stream(2), start=1): - if step == 1: - assert values == { - "inbox": 3, - } - # modify inbox value - values["inbox"] = 5 - elif step == 2: - # output is different now - assert values == { - "output": 6, - } - graph = Graph() graph.add_node("add_one", add_one) graph.add_node("add_one_more", add_one) @@ -229,25 +206,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert 0, f"{step}:{values}" assert step == 3 - for step, values in enumerate(gapp.stream(2), start=1): - if step == 1: - assert values == { - "add_one": 3, - } - # modify value before next step - values["add_one"] = 5 - elif step == 2: - assert values == { - "add_one_more": 6, - } - elif step == 3: - assert values == { - "__end__": 6, - } - else: - assert 0, "Should not get here" - assert step == 3 - def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4569c91b3..e7b85ee04 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -192,22 +192,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: } assert step == 2 - step = 0 - async for values in app.astream(2): - step += 1 - if step == 1: - assert values == { - "inbox": 3, - } - # modify inbox value - values["inbox"] = 5 - elif step == 2: - # output is different now - assert values == { - "output": 6, - } - assert step == 2 - graph = Graph() graph.add_node("add_one", add_one) graph.add_node("add_one_more", add_one) @@ -235,26 +219,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: } assert step == 3 - step = 0 - async for values in gapp.astream(2): - step += 1 - if step == 1: - assert values == { - "add_one": 3, - } - # modify value before running next step - values["add_one"] = 5 - elif step == 2: - # output is different now - assert values == { - "add_one_more": 6, - } - elif step == 3: - assert values == { - "__end__": 6, - } - assert step == 3 - async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1)