diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 7c630579e..ad0b650d2 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -252,6 +252,7 @@ class Pregel( *, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt: Optional[Sequence[str]] = None, ) -> Iterator[Union[dict[str, Any], Any]]: try: if config["recursion_limit"] < 1: @@ -267,6 +268,7 @@ class Pregel( input_keys = self.input else: validate_keys(input_keys, self.channels) + interrupt = interrupt or self.interrupt # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one @@ -277,13 +279,19 @@ class Pregel( self.channels, checkpoint ) as channels, get_executor_for_config(config) as executor: # map inputs to channel updates - _apply_writes( - checkpoint, - channels, - deque(w for c in input for w in map_input(input_keys, c)), - config, - 0, - ) + if input_writes := deque( + w for c in input for w in map_input(input_keys, c) + ): + # discard any unfinished tasks from previous checkpoint + _prepare_next_tasks(checkpoint, processes, channels) + # apply input writes + _apply_writes( + checkpoint, + channels, + input_writes, + config, + 0, + ) read = partial(_read_channel, channels) @@ -371,9 +379,7 @@ class Pregel( self.checkpointer.put(config, checkpoint) # interrupt if any channel written to is in interrupt list - if any( - chan for chan, _ in pending_writes if chan in self.interrupt - ): + if any(chan for chan, _ in pending_writes if chan in interrupt): break # save end of run checkpoint @@ -400,6 +406,7 @@ class Pregel( *, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt: Optional[Sequence[str]] = None, ) -> AsyncIterator[Union[dict[str, Any], Any]]: try: if config["recursion_limit"] < 1: @@ -424,6 +431,7 @@ class Pregel( input_keys = self.input else: validate_keys(input_keys, self.channels) + interrupt = interrupt or self.interrupt # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one @@ -434,13 +442,19 @@ class Pregel( # create channels from checkpoint async with AsyncChannelsManager(self.channels, checkpoint) as channels: # map inputs to channel updates - _apply_writes( - checkpoint, - channels, - deque([w async for c in input for w in map_input(input_keys, c)]), - config, - 0, - ) + if input_writes := deque( + [w async for c in input for w in map_input(input_keys, c)] + ): + # discard any unfinished tasks from previous checkpoint + _prepare_next_tasks(checkpoint, processes, channels) + # apply input writes + _apply_writes( + checkpoint, + channels, + input_writes, + config, + 0, + ) read = partial(_read_channel, channels) @@ -535,9 +549,7 @@ class Pregel( await self.checkpointer.aput(config, checkpoint) # interrupt if any channel written to is in interrupt list - if any( - chan for chan, _ in pending_writes if chan in self.interrupt - ): + if any(chan for chan, _ in pending_writes if chan in interrupt): break # save end of run checkpoint diff --git a/tests/test_pregel.py b/tests/test_pregel.py index ebc04044a..f9fef277b 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -245,6 +245,40 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 3 +def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + + memory = MemorySaver() + app = Pregel( + nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"] + ) + + # start execution, stop at inbox + assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None + + # inbox == 3 + checkpoint = memory.get({"configurable": {"thread_id": 1}}) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4 + + # start execution again, stop at inbox + assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None + + # inbox == 21 + checkpoint = memory.get({"configurable": {"thread_id": 1}}) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None + assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5 + + def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 286db9a3e..e829e2919 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -252,6 +252,40 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: 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) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + + memory = MemorySaver() + app = Pregel( + nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"] + ) + + # start execution, stop at inbox + assert await app.ainvoke(2, {"configurable": {"thread_id": 1}}) is None + + # inbox == 3 + checkpoint = await memory.aget({"configurable": {"thread_id": 1}}) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 4 + + # start execution again, stop at inbox + assert await app.ainvoke(20, {"configurable": {"thread_id": 1}}) is None + + # inbox == 21 + checkpoint = await memory.aget({"configurable": {"thread_id": 1}}) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None + assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5 + + async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")