diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index af4ba4d8d..f85b33ba3 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,10 +1,12 @@ -from typing import Any +from dataclasses import dataclass +from typing import Any, Literal INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" CONFIG_KEY_RESUMING = "__pregel_resuming" +CONFIG_KEY_TASK_ID = "__pregel_task_id" INTERRUPT = "__interrupt__" ERROR = "__error__" TASKS = "__pregel_tasks" @@ -16,6 +18,7 @@ RESERVED = { CONFIG_KEY_READ, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_RESUMING, + CONFIG_KEY_TASK_ID, INPUT, } TAG_HIDDEN = "langsmith:hidden" @@ -93,3 +96,9 @@ class Send: and self.node == value.node and self.arg == value.arg ) + + +@dataclass +class Interrupt: + when: Literal["before", "during", "after"] + value: Any = None diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index ea54cc55f..00dc36e58 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -1,4 +1,7 @@ +from typing import Any + from langgraph.checkpoint.base import EmptyChannelError +from langgraph.constants import Interrupt class GraphRecursionError(RecursionError): @@ -29,7 +32,15 @@ class InvalidUpdateError(Exception): class GraphInterrupt(Exception): """Raised when a subgraph is interrupted.""" - pass + def __init__(self, interrupts: list[Interrupt]) -> None: + super().__init__(interrupts) + + +class NodeInterrupt(GraphInterrupt): + """Raised by a node to interrupt execution.""" + + def __init__(self, value: Any) -> None: + super().__init__([Interrupt("during", value)]) class EmptyInputError(Exception): @@ -42,6 +53,7 @@ __all__ = [ "GraphRecursionError", "InvalidUpdateError", "GraphInterrupt", + "NodeInterrupt", "EmptyInputError", "EmptyChannelError", ] diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 7b2fdb5f9..d2f0790cc 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -70,8 +70,9 @@ from langgraph.constants import ( CONFIG_KEY_SEND, ERROR, INTERRUPT, + Interrupt, ) -from langgraph.errors import GraphRecursionError, InvalidUpdateError +from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( AsyncManagedValuesManager, ManagedValuesManager, @@ -82,6 +83,7 @@ from langgraph.pregel.algo import ( apply_writes, local_read, prepare_next_tasks, + should_interrupt, ) from langgraph.pregel.debug import ( map_debug_task_results, @@ -381,6 +383,7 @@ class Pregel( saved.metadata.get("step", -1) + 1 if saved else -1, for_execution=False, ) + return StateSnapshot( read_channels(channels, self.stream_channels_asis), tuple(t.name for t in next_tasks), @@ -550,7 +553,7 @@ class Pregel( saved = self.checkpointer.get_tuple(config) checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() checkpoint_previous_versions = ( - saved.checkpoint["channel_versions"] if saved else {} + saved.checkpoint["channel_versions"].copy() if saved else {} ) step = saved.metadata.get("step", -1) if saved else -1 # merge configurable fields with previous checkpoint config @@ -576,7 +579,7 @@ class Pregel( create_checkpoint(checkpoint, None, step), { "source": "update", - "step": step, + "step": step + 1, "writes": {}, }, {}, @@ -606,7 +609,11 @@ class Pregel( if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") # update channels - with ChannelsManager(self.channels, checkpoint, config) as channels: + with ChannelsManager( + self.channels, checkpoint, config + ) as channels, ManagedValuesManager( + self.managed_values_dict, ensure_config(config) + ) as managed: # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() if not writers: @@ -640,19 +647,43 @@ class Pregel( apply_writes( checkpoint, channels, [task], self.checkpointer.get_next_version ) - - new_versions = get_new_channel_versions( - checkpoint_previous_versions, checkpoint["channel_versions"] - ) + checkpoint = create_checkpoint(checkpoint, channels, step + 1) + # check interrupt before + if tasks := should_interrupt( + checkpoint, + self.interrupt_before_nodes, + prepare_next_tasks( + checkpoint, + self.nodes, + channels, + managed, + config, + step + 2, + for_execution=False, + ), + ): + for t in tasks: + self.checkpointer.put_writes( + { + "configurable": { + **checkpoint_config["configurable"], + "checkpoint_id": checkpoint["id"], + } + }, + [(INTERRUPT, Interrupt("before"))], + t.id, + ) return self.checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, step + 1), + checkpoint, { "source": "update", "step": step + 1, "writes": {as_node: values}, }, - new_versions, + get_new_channel_versions( + checkpoint_previous_versions, checkpoint["channel_versions"] + ), ) async def aupdate_state( @@ -668,7 +699,7 @@ class Pregel( saved = await self.checkpointer.aget_tuple(config) checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() checkpoint_previous_versions = ( - saved.checkpoint["channel_versions"] if saved else {} + saved.checkpoint["channel_versions"].copy() if saved else {} ) step = saved.metadata.get("step", -1) if saved else -1 # merge configurable fields with previous checkpoint config @@ -694,7 +725,7 @@ class Pregel( create_checkpoint(checkpoint, None, step), { "source": "update", - "step": step, + "step": step + 1, "writes": {}, }, {}, @@ -722,7 +753,11 @@ class Pregel( if as_node not in self.nodes: raise InvalidUpdateError(f"Node {as_node} does not exist") # update channels, acting as the chosen node - async with AsyncChannelsManager(self.channels, checkpoint, config) as channels: + async with AsyncChannelsManager( + self.channels, checkpoint, config + ) as channels, AsyncManagedValuesManager( + self.managed_values_dict, ensure_config(config) + ) as managed: # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() if not writers: @@ -756,19 +791,47 @@ class Pregel( apply_writes( checkpoint, channels, [task], self.checkpointer.get_next_version ) - - new_versions = get_new_channel_versions( - checkpoint_previous_versions, checkpoint["channel_versions"] - ) + checkpoint = create_checkpoint(checkpoint, channels, step + 1) + # check interrupt before + if tasks := should_interrupt( + checkpoint, + self.interrupt_before_nodes, + prepare_next_tasks( + checkpoint, + self.nodes, + channels, + managed, + config, + step + 2, + for_execution=False, + ), + ): + await asyncio.gather( + *( + self.checkpointer.aput_writes( + { + "configurable": { + **checkpoint_config["configurable"], + "checkpoint_id": checkpoint["id"], + } + }, + [(INTERRUPT, Interrupt("before"))], + t.id, + ) + for t in tasks + ) + ) return await self.checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, channels, step + 1), + checkpoint, { "source": "update", "step": step + 1, "writes": {as_node: values}, }, - new_versions, + get_new_channel_versions( + checkpoint_previous_versions, checkpoint["channel_versions"] + ), ) def _defaults( @@ -975,6 +1038,7 @@ class Pregel( for task in loop.tasks if not task.writes } + all_futures = futures.copy() end_time = ( self.step_timeout + time.monotonic() if self.step_timeout @@ -998,7 +1062,12 @@ class Pregel( task = futures.pop(fut) if exc := _exception(fut): # save error to checkpointer - loop.put_writes(task.id, [(ERROR, exc)]) + if isinstance(exc, GraphInterrupt): + loop.put_writes( + task.id, [(INTERRUPT, i) for i in exc.args[0]] + ) + else: + loop.put_writes(task.id, [(ERROR, exc)]) else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1026,7 +1095,7 @@ class Pregel( break # panic on failure or timeout - _panic_or_proceed(done, inflight, loop.step) + _panic_or_proceed(all_futures, loop.step) # don't keep futures around in memory longer than needed del done, inflight, futures # debug flag @@ -1223,6 +1292,7 @@ class Pregel( for task in loop.tasks if not task.writes } + all_futures = futures.copy() end_time = ( self.step_timeout + aioloop.time() if self.step_timeout @@ -1244,7 +1314,12 @@ class Pregel( task = futures.pop(fut) if exc := _exception(fut): # save error to checkpointer - loop.put_writes(task.id, [(ERROR, exc)]) + if isinstance(exc, GraphInterrupt): + loop.put_writes( + task.id, [(INTERRUPT, i) for i in exc.args[0]] + ) + else: + loop.put_writes(task.id, [(ERROR, exc)]) else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1274,7 +1349,7 @@ class Pregel( break # panic on failure or timeout - _panic_or_proceed(done, inflight, loop.step, asyncio.TimeoutError) + _panic_or_proceed(all_futures, loop.step, asyncio.TimeoutError) # don't keep futures around in memory longer than needed del done, inflight, futures # debug flag @@ -1421,9 +1496,8 @@ def _should_stop_others( for fut in done: if fut.cancelled(): return True - if fut.exception() is not None: - # TODO don't stop others if exception is interrupt - return True + if exc := fut.exception(): + return not isinstance(exc, GraphInterrupt) else: return False @@ -1441,14 +1515,20 @@ def _exception( def _panic_or_proceed( - done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], - inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], + futs: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], step: int, timeout_exc_cls: Type[Exception] = TimeoutError, ) -> None: + done: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set() + inflight: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set() + for fut in futs: + if fut.done(): + done.add(fut) + else: + inflight.add(fut) while done: # if any task failed - if exc := done.pop().exception(): + if exc := _exception(done.pop()): # cancel all pending tasks while inflight: inflight.pop().cancel() diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index fd725627a..5922fd14e 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -38,6 +38,7 @@ from langgraph.constants import ( CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, + CONFIG_KEY_TASK_ID, INTERRUPT, RESERVED, TAG_HIDDEN, @@ -68,26 +69,28 @@ def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], tasks: list[PregelExecutableTask], -) -> bool: +) -> list[PregelExecutableTask]: version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) null_version = version_type() seen = checkpoint["versions_seen"].get(INTERRUPT, {}) + # interrupt if any channel has been updated since last interrupt + any_updates_since_prev_interrupt = any( + version > seen.get(chan, null_version) + for chan, version in checkpoint["channel_versions"].items() + ) + # and any triggered node is in interrupt_nodes list return ( - # interrupt if any channel has been updated since last interrupt - any( - version > seen.get(chan, null_version) - for chan, version in checkpoint["channel_versions"].items() - ) - # and any triggered node is in interrupt_nodes list - and any( - task.name + [ + task for task in tasks if ( (not task.config or TAG_HIDDEN not in task.config.get("tags")) if interrupt_nodes == "*" else task.name in interrupt_nodes ) - ) + ] + if any_updates_since_prev_interrupt + else [] ) @@ -308,6 +311,7 @@ def prepare_next_tasks( else None ), configurable={ + CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe CONFIG_KEY_SEND: partial( local_write, writes.extend, processes, channels @@ -398,6 +402,7 @@ def prepare_next_tasks( else None ), configurable={ + CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe CONFIG_KEY_SEND: partial( local_write, writes.extend, processes, channels diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 931f3a180..5102cd8c5 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -10,7 +10,7 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite -from langgraph.constants import ERROR, TAG_HIDDEN +from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN from langgraph.pregel.io import read_channels from langgraph.pregel.types import PregelExecutableTask, PregelTask @@ -32,6 +32,7 @@ class CheckpointTask(TypedDict): id: str name: str error: Optional[str] + interrupts: list[dict] class CheckpointPayload(TypedDict): @@ -149,6 +150,7 @@ def map_debug_checkpoint( else { "id": t.id, "name": t.name, + "interrupts": t.interrupts, } for t in tasks_w_writes(tasks, pending_writes) ], @@ -190,8 +192,11 @@ def print_step_writes( def print_step_checkpoint( - step: int, channels: Mapping[str, BaseChannel], whitelist: Sequence[str] + metadata: CheckpointMetadata, + channels: Mapping[str, BaseChannel], + whitelist: Sequence[str], ) -> None: + step = metadata["step"] print( f"{get_colored_text(f'[{step}:checkpoint]', color='blue')} " + get_bolded_text(f"State at the end of step {step}:\n") @@ -203,6 +208,7 @@ def tasks_w_writes( tasks: list[PregelExecutableTask], pending_writes: Optional[list[PendingWrite]], ) -> tuple[PregelTask, ...]: + pending_writes = pending_writes or [] return tuple( PregelTask( task.id, @@ -210,12 +216,14 @@ def tasks_w_writes( next( ( exc - for tid, n, exc in pending_writes or [] - if tid == task.id - if n == ERROR + for tid, n, exc in pending_writes + if tid == task.id and n == ERROR ), None, ), + tuple( + v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT + ), ) for task in tasks ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f43aa25a7..f48f86fcd 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -45,6 +45,7 @@ from langgraph.constants import ( ERROR, INPUT, INTERRUPT, + Interrupt, ) from langgraph.errors import EmptyInputError, GraphInterrupt from langgraph.managed.base import ( @@ -206,10 +207,13 @@ class PregelLoop: } ) # after execution, check if we should interrupt - if should_interrupt(self.checkpoint, interrupt_after, self.tasks): + if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks): self.status = "interrupt_after" + interrupts = [(t.id, Interrupt("after")) for t in tasks] + for tid, interrupt in interrupts: + self.put_writes(tid, [(INTERRUPT, interrupt)]) if self.is_nested: - raise GraphInterrupt(self) + raise GraphInterrupt([i[1] for i in interrupts]) else: return False else: @@ -258,7 +262,7 @@ class PregelLoop: # if there are pending writes from a previous loop, apply them if self.checkpoint_pending_writes: for tid, k, v in self.checkpoint_pending_writes: - if k == ERROR: # TODO same for INTERRUPT + if k in (ERROR, INTERRUPT): continue if task := next((t for t in self.tasks if t.id == tid), None): task.writes.append((k, v)) @@ -273,10 +277,13 @@ class PregelLoop: ) # before execution, check if we should interrupt - if should_interrupt(self.checkpoint, interrupt_before, self.tasks): + if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" + interrupts = [(t.id, Interrupt("before")) for t in tasks] + for tid, interrupt in interrupts: + self.put_writes(tid, [(INTERRUPT, interrupt)]) if self.is_nested: - raise GraphInterrupt() + raise GraphInterrupt([i[1] for i in interrupts]) else: return False @@ -394,7 +401,7 @@ class PregelLoop: exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - if exc_type is GraphInterrupt and not self.is_nested: + if isinstance(exc_value, GraphInterrupt) and not self.is_nested: return True @@ -409,7 +416,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): ) -> None: super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) self.stack = ExitStack() - self.stack.push(self._suppress_interrupt) if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.put_writes @@ -457,6 +463,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): self.managed = self.stack.enter_context( ManagedValuesManager(self.graph.managed_values_dict, self.config) ) + self.stack.push(self._suppress_interrupt) self.status = "pending" self.step = self.checkpoint_metadata["step"] + 1 self.stop = self.step + self.config["recursion_limit"] + 1 @@ -486,7 +493,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): ) -> None: super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) self.stack = AsyncExitStack() - self.stack.push(self._suppress_interrupt) if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.aput_writes @@ -536,6 +542,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): self.managed = await self.stack.enter_async_context( AsyncManagedValuesManager(self.graph.managed_values_dict, self.config) ) + self.stack.push(self._suppress_interrupt) self.status = "pending" self.step = self.checkpoint_metadata["step"] + 1 self.stop = self.step + self.config["recursion_limit"] + 1 diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 8343329a4..a8c5a7beb 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -4,6 +4,7 @@ from typing import Any, Callable, Literal, NamedTuple, Optional, Type, Union from langchain_core.runnables import Runnable, RunnableConfig from langgraph.checkpoint.base import CheckpointMetadata +from langgraph.constants import Interrupt def default_retry_on(exc: Exception) -> bool: @@ -60,6 +61,7 @@ class PregelTask(NamedTuple): id: str name: str error: Optional[Exception] = None + interrupts: tuple[Interrupt, ...] = () class PregelExecutableTask(NamedTuple): diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 644c99b95..c92bc4688 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -682,6 +682,24 @@ ''' # --- +# name: test_dynamic_interrupt + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_in_one_fan_out_state_graph_waiting_edge ''' graph TD; diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 14a96b577..138d4dba5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -52,8 +52,8 @@ from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.sqlite import SqliteSaver -from langgraph.constants import ERROR, Send -from langgraph.errors import InvalidUpdateError +from langgraph.constants import ERROR, Interrupt, Send +from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph from langgraph.graph.graph import START from langgraph.graph.message import MessageGraph, add_messages @@ -2195,7 +2195,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2239,7 +2239,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "input": "what is weather in sf", }, }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2400,7 +2400,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2999,7 +2999,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3039,7 +3039,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3148,7 +3148,7 @@ def test_conditional_state_graph( values={ "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "agent"),), + tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3173,7 +3173,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3226,7 +3226,7 @@ def test_conditional_state_graph( ) ], }, - tasks=(PregelTask(AnyStr(), "agent"),), + tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4828,7 +4828,13 @@ def test_message_graph( id="ai1", ), ], - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=( + PregelTask( + AnyStr(), + "tools", + interrupts=(Interrupt("before"),), + ), + ), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4873,7 +4879,7 @@ def test_message_graph( ], ), ], - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4955,7 +4961,7 @@ def test_message_graph( id="ai2", ), ], - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5547,7 +5553,13 @@ def test_root_graph( id="ai1", ), ], - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=( + PregelTask( + AnyStr(), + "tools", + interrupts=(Interrupt("before"),), + ), + ), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5592,7 +5604,7 @@ def test_root_graph( ], ), ], - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5674,7 +5686,7 @@ def test_root_graph( id="ai2", ), ], - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -6120,6 +6132,72 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] +def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def tool_two_node(s: State) -> State: + if s["market"] == "DE": + raise NodeInterrupt("Just because...") + return {"my_key": " all good"} + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { + "my_key": "value", + "market": "DE", + } + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + with SqliteSaver.from_conn_string(":memory:") as saver: + tool_two = tool_two_graph.compile(checkpointer=saver) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("during", "Just because..."),), + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + def test_start_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] @@ -6172,7 +6250,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_slow", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6206,7 +6290,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6240,7 +6330,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6251,7 +6347,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6344,7 +6446,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"my_key": "value", "market": "DE"}, }, "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__"}], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, }, { @@ -6373,7 +6475,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": None, }, "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare"}], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, }, { @@ -6423,7 +6525,9 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"prepare": {"my_key": " prepared"}}, }, "next": ["tool_two_slow"], - "tasks": [{"id": AnyStr(), "name": "tool_two_slow"}], + "tasks": [ + {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} + ], }, }, { @@ -6473,7 +6577,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"tool_two_slow": {"my_key": " slow"}}, }, "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish"}], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, }, { @@ -6544,7 +6648,13 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_slow", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6582,7 +6692,13 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6629,7 +6745,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slow", "market": "DE", }, - tasks=(PregelTask(AnyStr(), "finish"),), + tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),), next=("finish",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6648,7 +6764,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slower", "market": "DE", }, - tasks=(PregelTask(AnyStr(), "finish"),), + tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),), next=("finish",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6905,7 +7021,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], }, - tasks=(PregelTask(AnyStr(), "qa"),), + tasks=(PregelTask(AnyStr(), "qa", interrupts=(Interrupt("before"),)),), next=("qa",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -7856,7 +7972,7 @@ def test_nested_graph_interrupts( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": "my value"}, config, debug=True) == { + assert app.invoke({"my_key": "my value"}, config) == { "my_key": "hi my value", } assert list(app.get_state_history(config)) == [ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b50ecd36a..7c5742aa0 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -47,8 +47,8 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver -from langgraph.constants import ERROR, Send -from langgraph.errors import InvalidUpdateError +from langgraph.constants import ERROR, Interrupt, Send +from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph, StateGraph from langgraph.graph.graph import START from langgraph.graph.message import MessageGraph, add_messages @@ -204,6 +204,61 @@ async def test_node_cancellation_on_other_node_exception() -> None: assert inner_task_cancelled +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_node_not_cancelled_on_other_node_interrupted( + checkpointer_name: str, request: pytest.FixtureRequest +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + hello: str + + awhiles = 0 + inner_task_cancelled = False + + async def awhile(input: State) -> None: + nonlocal awhiles + + awhiles += 1 + try: + await asyncio.sleep(1) + return {"hello": "again"} + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + async def iambad(input: State) -> None: + if input["hello"] != "bye": + raise NodeInterrupt("I am bad") + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.add_node("bad", iambad) + builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) + + graph = builder.compile(checkpointer=checkpointer) + thread = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} + + assert not inner_task_cancelled + assert awhiles == 1 + + assert await graph.ainvoke(None, thread, debug=True) is None + + assert not inner_task_cancelled + assert awhiles == 1 + + assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} + + assert not inner_task_cancelled + assert awhiles == 2 + + async def test_step_timeout_on_stream_hang() -> None: inner_task_cancelled = False @@ -2366,7 +2421,13 @@ async def test_conditional_graph() -> None: ), }, }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=( + PregelTask( + AnyStr(), + "tools", + interrupts=(Interrupt("before"),), + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2414,7 +2475,13 @@ async def test_conditional_graph() -> None: "input": "what is weather in sf", }, }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=( + PregelTask( + AnyStr(), + "tools", + interrupts=(Interrupt("before"),), + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2586,7 +2653,13 @@ async def test_conditional_graph() -> None: ), }, }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=( + PregelTask( + AnyStr(), + "tools", + interrupts=(Interrupt("before"),), + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3152,7 +3225,13 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=( + PregelTask( + AnyStr(), + "tools", + interrupts=(Interrupt("before"),), + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3196,7 +3275,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools"),), + tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4746,7 +4825,13 @@ async def test_start_branch_then() -> None: ] assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_slow", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -4788,7 +4873,13 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -4830,7 +4921,13 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -4845,7 +4942,13 @@ async def test_start_branch_then() -> None: await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -4945,12 +5048,7 @@ async def test_branch_then() -> None: "writes": {"my_key": "value", "market": "DE"}, }, "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - } - ], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, }, { @@ -4979,12 +5077,7 @@ async def test_branch_then() -> None: "writes": None, }, "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - } - ], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, }, { @@ -5035,10 +5128,7 @@ async def test_branch_then() -> None: }, "next": ["tool_two_slow"], "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - } + {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} ], }, }, @@ -5089,12 +5179,7 @@ async def test_branch_then() -> None: "writes": {"tool_two_slow": {"my_key": " slow"}}, }, "next": ["finish"], - "tasks": [ - { - "id": AnyStr(), - "name": "finish", - } - ], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, }, { @@ -5165,7 +5250,13 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_slow", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5211,7 +5302,13 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + tasks=( + PregelTask( + AnyStr(), + "tool_two_fast", + interrupts=(Interrupt("before"),), + ), + ), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[