From 89a085992848bd0a6daefab927e5d53c305c93b6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 4 Nov 2024 14:42:50 -0800 Subject: [PATCH 01/22] Execute Sends in same super step that triggered them - Keep old code path for compatibility with existing checkpoints - Keep a similar order of application of updates, in some cases there will be no visible change - Update task path for Sends to contain the path of all the parent tasks (multiple parents when a Send task creates another Send) - That lineage path is used to ensure order of application of updates respects their logical lineage (ie updates from parents always applied before their child tasks) - Move Interrupt writes to use negative indexes, which allow replacing/shadowing (when task is re-run it may interrupt again, or succeed) - Runner will now attempt to schedule new Send tasks as soon as the write is received (ie while the originating node is still running) - Update kafka scheduler to support new Send behavior --- .../langgraph/checkpoint/base/__init__.py | 3 +- .../langgraph/checkpoint/serde/types.py | 1 + libs/langgraph/langgraph/pregel/__init__.py | 63 +- libs/langgraph/langgraph/pregel/algo.py | 197 +++- libs/langgraph/langgraph/pregel/debug.py | 2 +- libs/langgraph/langgraph/pregel/loop.py | 122 ++- libs/langgraph/langgraph/pregel/retry.py | 20 +- libs/langgraph/langgraph/pregel/runner.py | 142 ++- libs/langgraph/langgraph/pregel/write.py | 6 +- libs/langgraph/langgraph/types.py | 4 +- libs/langgraph/tests/test_algo.py | 10 +- libs/langgraph/tests/test_pregel.py | 839 +++++++++++++----- libs/langgraph/tests/test_pregel_async.py | 800 +++++++++++++---- libs/scheduler-kafka/Makefile | 2 +- .../langgraph/scheduler/kafka/executor.py | 22 +- .../langgraph/scheduler/kafka/orchestrator.py | 24 +- .../langgraph/scheduler/kafka/types.py | 4 +- libs/scheduler-kafka/tests/test_push.py | 206 +++++ libs/scheduler-kafka/tests/test_push_sync.py | 208 +++++ 19 files changed, 2143 insertions(+), 532 deletions(-) create mode 100644 libs/scheduler-kafka/tests/test_push.py create mode 100644 libs/scheduler-kafka/tests/test_push_sync.py diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 93d510daa..a63bbce28 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -24,6 +24,7 @@ from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_ from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( ERROR, + INTERRUPT, SCHEDULED, ChannelProtocol, SendProtocol, @@ -449,4 +450,4 @@ Special writes (e.g. errors) map to negative indices, to avoid those writes from conflicting with regular writes. Each Checkpointer implementation should use this mapping in put_writes. """ -WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2} +WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2, INTERRUPT: -3} diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 154b1450b..9286e9b19 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -12,6 +12,7 @@ from typing_extensions import Self ERROR = "__error__" SCHEDULED = "__scheduled__" +INTERRUPT = "__interrupt__" TASKS = "__pregel_tasks" Value = TypeVar("Value", covariant=True) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index b491383df..76020cf89 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -69,6 +69,7 @@ from langgraph.constants import ( INTERRUPT, NS_END, NS_SEP, + PUSH, SCHEDULED, ) from langgraph.errors import ( @@ -98,7 +99,13 @@ from langgraph.pregel.utils import find_subgraph_pregel, get_new_channel_version from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, LoopProtocol, StateSnapshot, StreamMode +from langgraph.types import ( + All, + Checkpointer, + LoopProtocol, + StateSnapshot, + StreamMode, +) from langgraph.utils.config import ( ensure_config, merge_configs, @@ -468,6 +475,7 @@ class Pregel(PregelProtocol): # tasks for this checkpoint next_tasks = prepare_next_tasks( saved.checkpoint, + saved.pending_writes or [], self.nodes, channels, managed, @@ -570,6 +578,7 @@ class Pregel(PregelProtocol): # tasks for this checkpoint next_tasks = prepare_next_tasks( saved.checkpoint, + saved.pending_writes or [], self.nodes, channels, managed, @@ -922,6 +931,7 @@ class Pregel(PregelProtocol): # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, + saved.pending_writes, self.nodes, channels, managed, @@ -1001,8 +1011,14 @@ class Pregel(PregelProtocol): ), ) # save task writes - if saved: - checkpointer.put_writes(checkpoint_config, task.writes, task_id) + # channel writes are saved to current checkpoint + # push writes are saved to next checkpoint + channel_writes, push_writes = ( + [w for w in task.writes if w[0] != PUSH], + [w for w in task.writes if w[0] == PUSH], + ) + if saved and channel_writes: + checkpointer.put_writes(checkpoint_config, channel_writes, task_id) # apply to checkpoint and save mv_writes = apply_writes( checkpoint, channels, [task], checkpointer.get_next_version @@ -1023,6 +1039,8 @@ class Pregel(PregelProtocol): checkpoint_previous_versions, checkpoint["channel_versions"] ), ) + if push_writes: + checkpointer.put_writes(next_config, push_writes, task_id) return patch_checkpoint_map(next_config, saved.metadata if saved else None) async def aupdate_state( @@ -1132,6 +1150,7 @@ class Pregel(PregelProtocol): # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, + saved.pending_writes, self.nodes, channels, managed, @@ -1208,14 +1227,23 @@ class Pregel(PregelProtocol): ), ) # save task writes - if saved: - await checkpointer.aput_writes(checkpoint_config, writes, task_id) + # channel writes are saved to current checkpoint + # push writes are saved to next checkpoint + channel_writes, push_writes = ( + [w for w in task.writes if w[0] != PUSH], + [w for w in task.writes if w[0] == PUSH], + ) + if saved and channel_writes: + await checkpointer.aput_writes( + checkpoint_config, channel_writes, task_id + ) # apply to checkpoint and save mv_writes = apply_writes( checkpoint, channels, [task], checkpointer.get_next_version ) assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) + # save checkpoint, after applying writes next_config = await checkpointer.aput( checkpoint_config, checkpoint, @@ -1230,6 +1258,9 @@ class Pregel(PregelProtocol): checkpoint_previous_versions, checkpoint["channel_versions"] ), ) + # save push writes + if push_writes: + await checkpointer.aput_writes(next_config, push_writes, task_id) return patch_checkpoint_map(next_config, saved.metadata if saved else None) def _defaults( @@ -1432,12 +1463,16 @@ class Pregel(PregelProtocol): specs=self.channels, output_keys=output_keys, stream_keys=self.stream_channels_asis, + interrupt_before=interrupt_before_, + interrupt_after=interrupt_after_, + manager=run_manager, debug=debug, ) as loop: # create runner runner = PregelRunner( submit=loop.submit, put_writes=loop.put_writes, + schedule_task=loop.accept_push, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming @@ -1468,12 +1503,7 @@ class Pregel(PregelProtocol): # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps - while loop.tick( - input_keys=self.input_channels, - interrupt_before=interrupt_before_, - interrupt_after=interrupt_after_, - manager=run_manager, - ): + while loop.tick(input_keys=self.input_channels): for _ in runner.tick( loop.tasks.values(), timeout=self.step_timeout, @@ -1654,11 +1684,15 @@ class Pregel(PregelProtocol): specs=self.channels, output_keys=output_keys, stream_keys=self.stream_channels_asis, + interrupt_before=interrupt_before_, + interrupt_after=interrupt_after_, + manager=run_manager, ) as loop: # create runner runner = PregelRunner( submit=loop.submit, put_writes=loop.put_writes, + schedule_task=loop.accept_push, use_astream=do_stream is not None, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) @@ -1678,12 +1712,7 @@ class Pregel(PregelProtocol): # channel updates from step N are only visible in step N+1 # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps - while loop.tick( - input_keys=self.input_channels, - interrupt_before=interrupt_before_, - interrupt_after=interrupt_after_, - manager=run_manager, - ): + while loop.tick(input_keys=self.input_channels): async for _ in runner.atick( loop.tasks.values(), timeout=self.step_timeout, diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index af71294ae..6ae0253ac 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -25,6 +25,7 @@ from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, + PendingWrite, V, copy_checkpoint, ) @@ -68,7 +69,7 @@ class WritesProtocol(Protocol): Implemented by PregelTaskWrites and PregelExecutableTask.""" @property - def path(self) -> tuple[Union[str, int], ...]: ... + def path(self) -> tuple[Union[str, int, tuple], ...]: ... @property def name(self) -> str: ... @@ -84,7 +85,7 @@ class PregelTaskWrites(NamedTuple): """Simplest implementation of WritesProtocol, for usage with writes that don't originate from a runnable task, eg. graph input, update_state, etc.""" - path: tuple[Union[str, int], ...] + path: tuple[Union[str, int, tuple], ...] name: str writes: Sequence[tuple[str, Any]] triggers: Sequence[str] @@ -172,7 +173,7 @@ def local_write( """Function injected under CONFIG_KEY_SEND in task config, to write to channels. Validates writes and forwards them to `commit` function.""" for chan, value in writes: - if chan == TASKS: + if chan == PUSH: if not isinstance(value, Send): raise InvalidUpdateError(f"Expected Send, got {value}") if value.node not in process_keys: @@ -194,8 +195,10 @@ def apply_writes( """Apply writes from a set of tasks (usually the tasks from a Pregel step) to the checkpoint and channels, and return managed values writes to be applied externally.""" - # sort tasks on path - tasks = sorted(tasks, key=lambda t: t.path) + # sort tasks on path, to ensure deterministic order for update application + # any path parts after the 3rd are ignored for sorting + # (we use them for eg. task ids which aren't good for sorting) + tasks = sorted(tasks, key=lambda t: t.path[:3]) # update seen versions for task in tasks: @@ -237,8 +240,10 @@ def apply_writes( for chan, val in task.writes: if chan == NO_WRITES: pass - elif chan == TASKS: + elif chan == TASKS: # TODO: remove branch in 1.0 checkpoint["pending_sends"].append(val) + elif chan == PUSH: + pass elif chan in channels: pending_writes_by_channel[chan].append(val) else: @@ -277,6 +282,7 @@ def apply_writes( @overload def prepare_next_tasks( checkpoint: Checkpoint, + pending_writes: Sequence[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -293,6 +299,7 @@ def prepare_next_tasks( @overload def prepare_next_tasks( checkpoint: Checkpoint, + pending_writes: Sequence[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -308,6 +315,7 @@ def prepare_next_tasks( def prepare_next_tasks( checkpoint: Checkpoint, + pending_writes: Sequence[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -322,13 +330,14 @@ def prepare_next_tasks( """Prepare the set of tasks that will make up the next Pregel step. This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered by edges).""" - tasks: dict[str, Union[PregelTask, PregelExecutableTask]] = {} - # Consume pending packets - for idx, _ in enumerate(checkpoint["pending_sends"]): + tasks: list[Union[PregelTask, PregelExecutableTask]] = [] + # Consume pending_sends from previous step (legacy version of Send) + for idx, _ in enumerate(checkpoint["pending_sends"]): # TODO: remove branch in 1.0 if task := prepare_single_task( (PUSH, idx), None, checkpoint=checkpoint, + pending_writes=pending_writes, processes=processes, channels=channels, managed=managed, @@ -339,7 +348,7 @@ def prepare_next_tasks( checkpointer=checkpointer, manager=manager, ): - tasks[task.id] = task + tasks.append(task) # Check if any processes should be run in next step # If so, prepare the values to be passed to them for name in processes: @@ -347,6 +356,7 @@ def prepare_next_tasks( (PULL, name), None, checkpoint=checkpoint, + pending_writes=pending_writes, processes=processes, channels=channels, managed=managed, @@ -357,15 +367,74 @@ def prepare_next_tasks( checkpointer=checkpointer, manager=manager, ): - tasks[task.id] = task - return tasks + tasks.append(task) + # Consume pending Sends from this step (new version of Send) + if any(c == PUSH for _, c, _ in pending_writes): + # group writes by task id + grouped_by_task = defaultdict(list) + for tid, c, _ in pending_writes: + grouped_by_task[tid].append(c) + # prepare send tasks from grouped writes + # 1. start from sends originating from existing tasks + tidx = 0 + while tidx < len(tasks): + task = tasks[tidx] + if twrites := grouped_by_task.pop(task.id, None): + for idx, c in enumerate(twrites): + if c != PUSH: + continue + if next_task := prepare_single_task( + (PUSH, task.path, idx, task.id), + None, + checkpoint=checkpoint, + pending_writes=pending_writes, + processes=processes, + channels=channels, + managed=managed, + config=config, + step=step, + for_execution=for_execution, + store=store, + checkpointer=checkpointer, + manager=manager, + ): + tasks.append(next_task) + tidx += 1 + # key tasks by id + task_map = {t.id: t for t in tasks} + # 2. create new tasks for remaining sends (eg. from update_state) + for tid, writes in grouped_by_task.items(): + task = task_map.get(tid) + for idx, c in enumerate(writes): + if c != PUSH: + continue + if next_task := prepare_single_task( + (PUSH, task.path if task else (), idx, tid), + None, + checkpoint=checkpoint, + pending_writes=pending_writes, + processes=processes, + channels=channels, + managed=managed, + config=config, + step=step, + for_execution=for_execution, + store=store, + checkpointer=checkpointer, + manager=manager, + ): + task_map[next_task.id] = next_task + else: + task_map = {t.id: t for t in tasks} + return task_map def prepare_single_task( - task_path: tuple[str, Union[int, str]], + task_path: tuple[Union[str, int, tuple], ...], task_id_checksum: Optional[str], *, checkpoint: Checkpoint, + pending_writes: Sequence[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, @@ -383,31 +452,75 @@ def prepare_single_task( parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") if task_path[0] == PUSH: - idx = int(task_path[1]) - if idx >= len(checkpoint["pending_sends"]): - return - packet = checkpoint["pending_sends"][idx] - if not isinstance(packet, Send): - logger.warning( - f"Ignoring invalid packet type {type(packet)} in pending sends" + if len(task_path) == 2: # TODO: remove branch in 1.0 + # legacy SEND tasks, executed in superstep n+1 + # (PUSH, idx of pending send) + idx = cast(int, task_path[1]) + if idx >= len(checkpoint["pending_sends"]): + return + packet = checkpoint["pending_sends"][idx] + if not isinstance(packet, Send): + logger.warning( + f"Ignoring invalid packet type {type(packet)} in pending sends" + ) + return + if packet.node not in processes: + logger.warning( + f"Ignoring unknown node name {packet.node} in pending sends" + ) + return + # create task id + triggers = [PUSH] + checkpoint_ns = ( + f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node ) + task_id = _uuid5_str( + checkpoint_id, + checkpoint_ns, + str(step), + packet.node, + PUSH, + str(idx), + ) + elif len(task_path) == 4: + # new PUSH tasks, executed in superstep n + # (PUSH, parent task path, idx of PUSH write, id of parent task) + task_path_t = cast(tuple[str, tuple, int, str], task_path) + writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]] + if task_path_t[2] >= len(writes_for_path): + logger.warning( + f"Ignoring invalid write index {task_path[2]} in pending writes" + ) + return + packet = writes_for_path[task_path_t[2]][2] + if not isinstance(packet, Send): + print("packet", task_path_t, writes_for_path) + logger.warning( + f"Ignoring invalid packet type {type(packet)} in pending writes" + ) + return + if packet.node not in processes: + logger.warning( + f"Ignoring unknown node name {packet.node} in pending writes" + ) + return + # create task id + triggers = [PUSH] + checkpoint_ns = ( + f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node + ) + task_id = _uuid5_str( + checkpoint_id, + checkpoint_ns, + str(step), + packet.node, + PUSH, + _tuple_str(task_path[1]), + str(task_path[2]), + ) + else: + logger.warning(f"Ignoring invalid PUSH task path {task_path}") return - if packet.node not in processes: - logger.warning(f"Ignoring unknown node name {packet.node} in pending sends") - return - # create task id - triggers = [PUSH] - checkpoint_ns = ( - f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node - ) - task_id = _uuid5_str( - checkpoint_id, - checkpoint_ns, - str(step), - packet.node, - PUSH, - str(idx), - ) task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" metadata = { "langgraph_step": step, @@ -417,7 +530,7 @@ def prepare_single_task( "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: - assert task_id == task_id_checksum + assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: proc = processes[packet.node] if node := proc.node: @@ -481,6 +594,7 @@ def prepare_single_task( else: return PregelTask(task_id, packet.node, task_path) elif task_path[0] == PULL: + # (PULL, node name) name = cast(str, task_path[1]) if name not in processes: return @@ -642,3 +756,12 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str: sha.update(b"".join(p.encode() for p in parts)) hex = sha.hexdigest() return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" + + +def _tuple_str(tup: Union[str, int, tuple]) -> str: + """Generate a string representation of a tuple.""" + return ( + f"({', '.join(_tuple_str(x) for x in tup)})" + if isinstance(tup, (tuple, list)) + else str(tup) + ) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index d772e7cba..4aac70c2e 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -208,7 +208,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: print( f"{get_colored_text(f'[{step}:tasks]', color='blue')} " + get_bolded_text( - f"Starting step {step} with {n_tasks} task{'s' if n_tasks != 1 else ''}:\n" + f"Starting {n_tasks} task{'s' if n_tasks != 1 else ''} for step {step}:\n" ) + "\n".join( f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}" diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a74b2daa9..2e10a7a84 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -52,6 +52,7 @@ from langgraph.constants import ( INPUT, INTERRUPT, NS_SEP, + PUSH, SCHEDULED, TAG_HIDDEN, ) @@ -74,6 +75,7 @@ from langgraph.pregel.algo import ( apply_writes, increment, prepare_next_tasks, + prepare_single_task, should_interrupt, ) from langgraph.pregel.debug import ( @@ -130,6 +132,9 @@ class PregelLoop(LoopProtocol): stream_keys: Union[str, Sequence[str]] skip_done_tasks: bool is_nested: bool + manager: Union[None, AsyncParentRunManager, ParentRunManager] + interrupt_after: Union[All, Sequence[str]] + interrupt_before: Union[All, Sequence[str]] checkpointer_get_next_version: GetNextVersion checkpointer_put_writes: Optional[ @@ -162,6 +167,7 @@ class PregelLoop(LoopProtocol): "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" ] tasks: dict[str, PregelExecutableTask] + to_interrupt: list[PregelExecutableTask] output: Union[None, dict[str, Any], Any] = None # public @@ -178,6 +184,9 @@ class PregelLoop(LoopProtocol): specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], output_keys: Union[str, Sequence[str]], stream_keys: Union[str, Sequence[str]], + interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, + interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, + manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, check_subgraphs: bool = True, debug: bool = False, ) -> None: @@ -194,6 +203,9 @@ class PregelLoop(LoopProtocol): self.specs = specs self.output_keys = output_keys self.stream_keys = stream_keys + self.interrupt_after = interrupt_after + self.interrupt_before = interrupt_before + self.manager = manager self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {}) self.skip_done_tasks = ( CONFIG_KEY_CHECKPOINT_ID not in config[CONF] @@ -263,13 +275,57 @@ class PregelLoop(LoopProtocol): # output writes self._output_writes(task_id, writes) + def accept_push( + self, task: PregelExecutableTask, write_idx: int + ) -> Optional[PregelExecutableTask]: + """Accept a PUSH from a task, potentially returning a new task to start.""" + # don't start if an earlier PUSH has already triggered an interrupt + if self.to_interrupt: + return + # don't start if we should interrupt *after* the original task + if should_interrupt(self.checkpoint, self.interrupt_after, [task]): + self.to_interrupt.append(task) + return + if pushed := cast( + Optional[PregelExecutableTask], + prepare_single_task( + (PUSH, task.path, write_idx, task.id), + None, + checkpoint=self.checkpoint, + pending_writes=[(task.id, *w) for w in task.writes], + processes=self.nodes, + channels=self.channels, + managed=self.managed, + config=self.config, + step=self.step, + for_execution=True, + store=self.store, + checkpointer=self.checkpointer, + manager=self.manager, + ), + ): + # don't start if we should interrupt *before* the new task + if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]): + self.to_interrupt.append(pushed) + return + # produce debug output + self._emit("debug", map_debug_tasks, self.step, [pushed]) + # debug flag + if self.debug: + print_step_tasks(self.step, [pushed]) + # save the new task + self.tasks[pushed.id] = pushed + # match any pending writes to the new task + if self.skip_done_tasks: + self._match_writes({pushed.id: pushed}) + # return the new task, to be started, if not run before + if not pushed.writes: + return pushed + def tick( self, *, input_keys: Union[str, Sequence[str]], - interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, - interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, - manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, ) -> bool: """Execute a single iteration of the Pregel loop. Returns True if more iterations are needed.""" @@ -278,6 +334,10 @@ class PregelLoop(LoopProtocol): if self.input not in (INPUT_DONE, INPUT_RESUMING): self._first(input_keys=input_keys) + elif self.to_interrupt: + # if we need to interrupt, do so + self.status = "interrupt_before" + raise GraphInterrupt() elif all(task.writes for task in self.tasks.values()): writes = [w for t in self.tasks.values() for w in t.writes] # debug flag @@ -322,7 +382,9 @@ class PregelLoop(LoopProtocol): } ) # after execution, check if we should interrupt - if should_interrupt(self.checkpoint, interrupt_after, self.tasks.values()): + if should_interrupt( + self.checkpoint, self.interrupt_after, self.tasks.values() + ): self.status = "interrupt_after" raise GraphInterrupt() else: @@ -336,16 +398,18 @@ class PregelLoop(LoopProtocol): # prepare next tasks self.tasks = prepare_next_tasks( self.checkpoint, + self.checkpoint_pending_writes, self.nodes, self.channels, self.managed, self.config, self.step, for_execution=True, - manager=manager, + manager=self.manager, store=self.store, checkpointer=self.checkpointer, ) + self.to_interrupt = [] # produce debug output if self._checkpointer_put_after_previous is not None: @@ -387,15 +451,12 @@ class PregelLoop(LoopProtocol): # if all tasks have finished, re-tick if all(task.writes for task in self.tasks.values()): - return self.tick( - input_keys=input_keys, - interrupt_after=interrupt_after, - interrupt_before=interrupt_before, - manager=manager, - ) + return self.tick(input_keys=input_keys) # before execution, check if we should interrupt - if should_interrupt(self.checkpoint, interrupt_before, self.tasks.values()): + if should_interrupt( + self.checkpoint, self.interrupt_before, self.tasks.values() + ): self.status = "interrupt_before" raise GraphInterrupt() @@ -464,6 +525,7 @@ class PregelLoop(LoopProtocol): # discard any unfinished tasks from previous checkpoint discard_tasks = prepare_next_tasks( self.checkpoint, + self.checkpoint_pending_writes, self.nodes, self.channels, self.managed, @@ -577,11 +639,33 @@ class PregelLoop(LoopProtocol): # save final output self.output = read_channels(self.channels, self.output_keys) if suppress: - # suppress interrupt + # emit one last "values" event, with pending writes applied + if ( + hasattr(self, "tasks") + and self.checkpoint_pending_writes + and any(task.writes for task in self.tasks.values()) + ): + mv_writes = apply_writes( + self.checkpoint, + self.channels, + self.tasks.values(), + self.checkpointer_get_next_version, + ) + for key, values in mv_writes.items(): + self._update_mv(key, values) + self._emit( + "values", + map_output_values, + self.output_keys, + [w for t in self.tasks.values() for w in t.writes], + self.channels, + ) + # emit INTERRUPT event self._emit( "updates", lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]), ) + # suppress interrupt return True def _emit( @@ -635,6 +719,9 @@ class SyncPregelLoop(PregelLoop, ContextManager): checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, + interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, + interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, check_subgraphs: bool = True, @@ -650,7 +737,10 @@ class SyncPregelLoop(PregelLoop, ContextManager): specs=specs, output_keys=output_keys, stream_keys=stream_keys, + interrupt_after=interrupt_after, + interrupt_before=interrupt_before, check_subgraphs=check_subgraphs, + manager=manager, debug=debug, ) self.stack = ExitStack() @@ -761,6 +851,9 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): checkpointer: Optional[BaseCheckpointSaver], nodes: Mapping[str, PregelNode], specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], + interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, + interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, + manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, check_subgraphs: bool = True, @@ -776,7 +869,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): specs=specs, output_keys=output_keys, stream_keys=stream_keys, + interrupt_after=interrupt_after, + interrupt_before=interrupt_before, check_subgraphs=check_subgraphs, + manager=manager, debug=debug, ) self.stack = AsyncExitStack() diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 60057493d..ea9162dc2 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -2,9 +2,15 @@ import asyncio import logging import random import time -from typing import Optional, Sequence +from functools import partial +from typing import Any, Callable, Optional, Sequence -from langgraph.constants import CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING +from langgraph.constants import ( + CONF, + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_RESUMING, + CONFIG_KEY_SEND, +) from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphInterrupt from langgraph.types import PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable @@ -15,12 +21,17 @@ logger = logging.getLogger(__name__) def run_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], + writer: Optional[ + Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None] + ] = None, ) -> None: """Run a task with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config + if writer is not None: + config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)}) while True: try: # clear any writes from previous attempts @@ -84,12 +95,17 @@ async def arun_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], stream: bool = False, + writer: Optional[ + Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None] + ] = None, ) -> None: """Run a task asynchronously with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config + if writer is not None: + config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)}) while True: try: # clear any writes from previous attempts diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index a2ca1c6cb..64e5c8d3c 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -14,7 +14,15 @@ from typing import ( cast, ) -from langgraph.constants import ERROR, INTERRUPT, NO_WRITES, TAG_HIDDEN +from langgraph.constants import ( + CONF, + CONFIG_KEY_SEND, + ERROR, + INTERRUPT, + NO_WRITES, + PUSH, + TAG_HIDDEN, +) from langgraph.errors import GraphDelegate, GraphInterrupt from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry @@ -31,6 +39,9 @@ class PregelRunner: *, submit: Submit, put_writes: Callable[[str, Sequence[tuple[str, Any]]], None], + schedule_task: Callable[ + [PregelExecutableTask, int], Optional[PregelExecutableTask] + ], use_astream: bool = False, node_finished: Optional[Callable[[str], None]] = None, ) -> None: @@ -38,6 +49,7 @@ class PregelRunner: self.put_writes = put_writes self.use_astream = use_astream self.node_finished = node_finished + self.schedule_task = schedule_task def tick( self, @@ -48,27 +60,58 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: + def writer( + task: PregelExecutableTask, writes: Sequence[tuple[str, Any]] + ) -> None: + prev_length = len(task.writes) + # delegate to the underlying writer + task.config[CONF][CONFIG_KEY_SEND](writes) + for idx, w in enumerate(task.writes): + # find the index for the newly inserted writes + if idx < prev_length: + continue + assert writes[idx - prev_length] is w + # bail if not a PUSH write + if w[0] != PUSH: + continue + # schedule the next task, if the callback returns one + if next_task := self.schedule_task(task, idx): + # if the parent task was retried, + # the next task might already be running + if any( + t == next_task.id for t in futures.values() if t is not None + ): + continue + # schedule the next task + futures[ + self.submit( + run_with_retry, + next_task, + retry_policy, + writer=writer, + __reraise_on_exit__=reraise, + ) + ] = next_task + tasks = tuple(tasks) + futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {} # give control back to the caller yield # fast path if single task with no timeout and no waiter if len(tasks) == 1 and timeout is None and get_waiter is None: t = tasks[0] try: - run_with_retry(t, retry_policy) + run_with_retry(t, retry_policy, writer=writer) self.commit(t, None) except Exception as exc: self.commit(t, exc) if reraise: raise - return + if not futures: # maybe `t` schuduled another task + return # add waiter task if requested if get_waiter is not None: - futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = { - get_waiter(): None - } - else: - futures = {} + futures[get_waiter()] = None # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes @@ -79,10 +122,11 @@ class PregelRunner: run_with_retry, t, retry_policy, + writer=writer, __reraise_on_exit__=reraise, ) ] = t - all_futures = futures.copy() + done_futures: set[concurrent.futures.Future] = set() end_time = timeout + time.monotonic() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = concurrent.futures.wait( @@ -99,6 +143,8 @@ class PregelRunner: if inflight and get_waiter is not None: futures[get_waiter()] = None else: + # store for panic check + done_futures.add(fut) # task finished, commit writes self.commit(task, _exception(fut)) else: @@ -110,7 +156,10 @@ class PregelRunner: # give control back to the caller yield # panic on failure or timeout - _panic_or_proceed(all_futures, panic=reraise) + _panic_or_proceed( + done_futures.union(f for f, t in futures.items() if t is not None), + panic=reraise, + ) async def atick( self, @@ -121,28 +170,67 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: + def writer( + task: PregelExecutableTask, writes: Sequence[tuple[str, Any]] + ) -> None: + prev_length = len(task.writes) + # delegate to the underlying writer + task.config[CONF][CONFIG_KEY_SEND](writes) + for idx, w in enumerate(task.writes): + # find the index for the newly inserted writes + if idx < prev_length: + continue + assert writes[idx - prev_length] is w + # bail if not a PUSH write + if w[0] != PUSH: + continue + # schedule the next task, if the callback returns one + if next_task := self.schedule_task(task, idx): + # if the parent task was retried, + # the next task might already be running + if any( + t == next_task.id for t in futures.values() if t is not None + ): + continue + # schedule the next task + futures[ + cast( + asyncio.Future, + self.submit( + arun_with_retry, + next_task, + retry_policy, + stream=self.use_astream, + writer=writer, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), + ) + ] = next_task + loop = asyncio.get_event_loop() tasks = tuple(tasks) + futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {} # give control back to the caller yield # fast path if single task with no waiter and no timeout if len(tasks) == 1 and get_waiter is None and timeout is None: t = tasks[0] try: - await arun_with_retry(t, retry_policy, stream=self.use_astream) + await arun_with_retry( + t, retry_policy, stream=self.use_astream, writer=writer + ) self.commit(t, None) except Exception as exc: self.commit(t, exc) if reraise: raise - return + if not futures: # maybe `t` schuduled another task + return # add waiter task if requested if get_waiter is not None: - futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = { - get_waiter(): None - } - else: - futures = {} + futures[get_waiter()] = None # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes @@ -156,13 +244,14 @@ class PregelRunner: t, retry_policy, stream=self.use_astream, + writer=writer, __name__=t.name, __cancel_on_exit__=True, __reraise_on_exit__=reraise, ), ) ] = t - all_futures = futures.copy() + done_futures: set[asyncio.Future] = set() end_time = timeout + loop.time() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = await asyncio.wait( @@ -179,6 +268,8 @@ class PregelRunner: if inflight and get_waiter is not None: futures[get_waiter()] = None else: + # store for panic check + done_futures.add(fut) # task finished, commit writes self.commit(task, _exception(fut)) else: @@ -194,7 +285,9 @@ class PregelRunner: fut.cancel() # panic on failure or timeout _panic_or_proceed( - all_futures, timeout_exc_cls=asyncio.TimeoutError, panic=reraise + done_futures.union(f for f, t in futures.items() if t is not None), + timeout_exc_cls=asyncio.TimeoutError, + panic=reraise, ) def commit( @@ -250,10 +343,7 @@ def _exception( def _panic_or_proceed( - futs: Union[ - dict[concurrent.futures.Future, Optional[PregelExecutableTask]], - dict[asyncio.Future, Optional[PregelExecutableTask]], - ], + futs: Union[set[concurrent.futures.Future], set[asyncio.Future]], *, timeout_exc_cls: Type[Exception] = TimeoutError, panic: bool = True, @@ -261,10 +351,8 @@ def _panic_or_proceed( """Cancel remaining tasks if any failed, re-raise exception if panic is True.""" done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() - for fut, val in futs.items(): - if val is None: - continue - elif fut.done(): + for fut in futs: + if fut.done(): done.add(fut) else: inflight.add(fut) diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 9975c7e5b..ba783453c 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -14,7 +14,7 @@ from typing import ( from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec -from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send +from langgraph.constants import CONF, CONFIG_KEY_SEND, PUSH, TASKS, Send from langgraph.errors import InvalidUpdateError from langgraph.utils.runnable import RunnableCallable @@ -112,14 +112,14 @@ class ChannelWrite(RunnableCallable): # validate for w in writes: if isinstance(w, ChannelWriteEntry): - if w.channel == TASKS: + if w.channel in (TASKS, PUSH): raise InvalidUpdateError( "Cannot write to the reserved channel TASKS" ) if w.value is PASSTHROUGH: raise InvalidUpdateError("PASSTHROUGH value must be replaced") # split packets and entries - sends = [(TASKS, packet) for packet in writes if isinstance(packet, Send)] + sends = [(PUSH, packet) for packet in writes if isinstance(packet, Send)] entries = [write for write in writes if isinstance(write, ChannelWriteEntry)] # process entries into values values = [ diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index b76f71d70..23ea2d881 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -110,7 +110,7 @@ class Interrupt: class PregelTask(NamedTuple): id: str name: str - path: tuple[Union[str, int], ...] + path: tuple[Union[str, int, tuple], ...] error: Optional[Exception] = None interrupts: tuple[Interrupt, ...] = () state: Union[None, RunnableConfig, "StateSnapshot"] = None @@ -127,7 +127,7 @@ class PregelExecutableTask(NamedTuple): retry_policy: Optional[RetryPolicy] cache_policy: Optional[CachePolicy] id: str - path: tuple[Union[str, int], ...] + path: tuple[Union[str, int, tuple], ...] scheduled: bool = False diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index 4e259f29e..9d6ec5942 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -11,13 +11,21 @@ def test_prepare_next_tasks() -> None: with ChannelsManager({}, checkpoint, config) as (channels, managed): assert ( prepare_next_tasks( - checkpoint, processes, channels, managed, config, 0, for_execution=False + checkpoint, + {}, + processes, + channels, + managed, + config, + 0, + for_execution=False, ) == {} ) assert ( prepare_next_tasks( checkpoint, + {}, processes, channels, managed, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 20b12beb2..bf92f489c 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1785,12 +1785,12 @@ def test_concurrent_emit_sends() -> None: "0", "1", "1.1", - "3.1", "2|1", "2|2", "2|3", "2|4", "3", + "3.1", ] @@ -1803,7 +1803,7 @@ def test_send_sequences() -> None: def __call__(self, state): update = ( [self.name] - if isinstance(state, list) # or isinstance(state, Control) + if isinstance(state, list) else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): @@ -1844,6 +1844,338 @@ def test_send_sequences() -> None: ] +@pytest.mark.repeat(20) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_dedupe_on_resume( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + if checkpointer_name == "duckdb": + pytest.skip("DuckDB isn't returning the right history") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + setattr(self, "__name__", name) + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Control): + state.state = update + return state + else: + return update + + def send_for_fun(state): + return [ + Send("2", Control(send=Send("2", 3))), + Send("2", Control(send=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + graph = builder.compile(checkpointer=checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke(["0"], thread1, debug=1) == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + ] + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # check state + state = graph.get_state(thread1) + assert state.next == ("flaky",) + # check history + history = [c for c in graph.get_state_history(thread1)] + assert len(history) == 2 + # resume execution + assert graph.invoke(None, thread1, debug=1) == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check state + state = graph.get_state(thread1) + assert state.next == () + # check history + history = [c for c in graph.get_state_history(thread1)] + assert ( + history[1] + == [ + StateSnapshot( + values=[ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"3": ["3"], "3.1": ["3.1"]}, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + ], + next=("3", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "1": ["1"], + "2": [ + ["2|Control(send=Send(node='2', arg=3))"], + ["2|Control(send=Send(node='flaky', arg=4))"], + ["2|3"], + ], + "flaky": ["flaky|4"], + }, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + ), + StateSnapshot( + values=["0"], + next=("1", "2", "2", "2", "flaky"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Control(send=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Control(send=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(Interrupt(value="Bahh", when="during"),), + state=None, + result=["flaky|4"], + ), + ), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": ["0"]}, + "thread_id": "1", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + ), + ][1] + ) + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_react_interrupt( request: pytest.FixtureRequest, checkpointer_name: str @@ -1996,23 +2328,9 @@ def test_send_react_interrupt( } }, metadata={ - "step": 1, + "step": 0, "source": "loop", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, + "writes": None, "parents": {}, "thread_id": "2", }, @@ -2025,10 +2343,34 @@ def test_send_react_interrupt( } }, tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), error=None, interrupts=(), state=None, @@ -2062,7 +2404,7 @@ def test_send_react_interrupt( } }, metadata={ - "step": 2, + "step": 1, "source": "update", "writes": { "agent": { @@ -2145,23 +2487,9 @@ def test_send_react_interrupt( } }, metadata={ - "step": 1, + "step": 0, "source": "loop", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, + "writes": None, "parents": {}, "thread_id": "3", }, @@ -2174,10 +2502,32 @@ def test_send_react_interrupt( } }, tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), error=None, interrupts=(), state=None, @@ -2232,7 +2582,7 @@ def test_send_react_interrupt( } }, metadata={ - "step": 2, + "step": 1, "source": "update", "writes": { "agent": { @@ -2264,7 +2614,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", (), 0, AnyStr()), error=None, interrupts=(), state=None, @@ -2443,23 +2793,9 @@ def test_send_react_interrupt_control( } }, metadata={ - "step": 1, + "step": 0, "source": "loop", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, + "writes": None, "parents": {}, "thread_id": "2", }, @@ -2472,10 +2808,34 @@ def test_send_react_interrupt_control( } }, tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), error=None, interrupts=(), state=None, @@ -2509,7 +2869,7 @@ def test_send_react_interrupt_control( } }, metadata={ - "step": 2, + "step": 1, "source": "update", "writes": { "agent": { @@ -5525,29 +5885,43 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], metadata={ "parents": {}, "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, + "step": 0, + "writes": None, "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, @@ -5578,14 +5952,14 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 2, + "step": 1, "writes": { "agent": { "messages": AIMessage( @@ -5679,8 +6053,40 @@ def test_state_graph_packets( ] }, tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -5688,25 +6094,14 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "loop", - "step": 4, + "step": 2, "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), }, }, "thread_id": "1", @@ -5753,7 +6148,7 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "update", - "step": 5, + "step": 3, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -5815,29 +6210,41 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], metadata={ "parents": {}, "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, + "step": 0, + "writes": None, "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, @@ -5868,14 +6275,14 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 2, + "step": 1, "writes": { "agent": { "messages": AIMessage( @@ -5969,8 +6376,40 @@ def test_state_graph_packets( ] }, tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -5978,25 +6417,14 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "loop", - "step": 4, + "step": 2, "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), }, }, "thread_id": "2", @@ -6043,7 +6471,7 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "update", - "step": 5, + "step": 3, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -10045,7 +10473,7 @@ def test_nested_graph_interrupts_parallel( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert app.invoke({"my_key": ""}, config, debug=True) == { - "my_key": "", + "my_key": " and parallel", } assert app.invoke(None, config, debug=True) == { @@ -10073,6 +10501,7 @@ def test_nested_graph_interrupts_parallel( config = {"configurable": {"thread_id": "3"}} assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -10089,6 +10518,7 @@ def test_nested_graph_interrupts_parallel( # while we're waiting for the node w/ interrupt inside to finish assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -10100,7 +10530,8 @@ def test_nested_graph_interrupts_parallel( app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ - {"my_key": ""} + {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -11838,10 +12269,19 @@ def test_send_to_nested_graphs( assert outer_state == StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), PregelTask( AnyStr(), "generate_joke", - (PUSH, 0), + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -11852,7 +12292,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, 1), + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -11871,22 +12311,16 @@ def test_send_to_nested_graphs( }, metadata={ "parents": {}, - "source": "loop", - "writes": None, - "step": 0, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, "thread_id": "1", }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, + parent_config=None, ) # check state of each of the inner tasks - assert graph.get_state(outer_state.tasks[0].state) == StateSnapshot( + assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( values={"subject": "cats - hohoho", "jokes": []}, next=("generate",), config={ @@ -11911,8 +12345,8 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), "langgraph_checkpoint_ns": AnyStr("generate_joke:"), "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, 0], - "langgraph_step": 1, + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1, AnyStr()], + "langgraph_step": 0, "langgraph_triggers": [PUSH], }, created_at=AnyStr(), @@ -11931,7 +12365,7 @@ def test_send_to_nested_graphs( }, tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), ) - assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( + assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot( values={"subject": "dogs - hohoho", "jokes": []}, next=("generate",), config={ @@ -11956,8 +12390,8 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), "langgraph_checkpoint_ns": AnyStr("generate_joke:"), "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, 1], - "langgraph_step": 1, + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2, AnyStr()], + "langgraph_step": 0, "langgraph_triggers": [PUSH], }, created_at=AnyStr(), @@ -11977,7 +12411,7 @@ def test_send_to_nested_graphs( tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), ) # update state of dogs joke graph - graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) + graph.update_state(outer_state.tasks[2].state, {"subject": "turtles - hohoho"}) # continue past interrupt assert sorted( @@ -12011,7 +12445,7 @@ def test_send_to_nested_graphs( {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, "thread_id": "1", }, created_at=AnyStr(), @@ -12053,58 +12487,6 @@ def test_send_to_nested_graphs( {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - tasks=( - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, 0), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about cats - hohoho"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about turtles - hohoho"]}, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, "step": 0, "thread_id": "1", }, @@ -12121,13 +12503,40 @@ def test_send_to_nested_graphs( values={"jokes": []}, tasks=( PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, result={"subjects": ["cats", "dogs"]}, ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about cats - hohoho"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about turtles - hohoho"]}, + ), ), - next=("__start__",), + next=("__start__", "generate_joke", "generate_joke"), config={ "configurable": { "thread_id": "1", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 66d3c2b8f..b32100a2e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -399,12 +399,13 @@ async def test_node_not_cancelled_on_other_node_interrupted( graph = builder.compile(checkpointer=checkpointer) thread = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} + # writes from "awhile" are applied to last chunk + assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "again"} assert not inner_task_cancelled assert awhiles == 1 - assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"} + assert await graph.ainvoke(None, thread, debug=True) == {"hello": "again"} assert not inner_task_cancelled assert awhiles == 1 @@ -2025,12 +2026,12 @@ async def test_concurrent_emit_sends() -> None: "0", "1", "1.1", - "3.1", "2|1", "2|2", "2|3", "2|4", "3", + "3.1", ] @@ -2074,8 +2075,8 @@ async def test_send_sequences() -> None: "0", "1", "3.1", - "2|Command(send=Send(node='2', arg=3))", - "2|Command(send=Send(node='2', arg=4))", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='2', arg=4))", "3", "2|3", "2|4", @@ -2083,6 +2084,324 @@ async def test_send_sequences() -> None: ] +@pytest.mark.repeat(20) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: + if checkpointer_name == "duckdb_aio": + pytest.skip("DuckDB isn't returning the right history") + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + setattr(self, "__name__", name) + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Control): + state.state = update + return state + else: + return update + + def send_for_fun(state): + return [ + Send("2", Control(send=Send("2", 3))), + Send("2", Control(send=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1, debug=1) == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + ] + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # resume execution + assert await graph.ainvoke(None, thread1, debug=1) == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=[ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"3": ["3"], "3.1": ["3.1"]}, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + ], + next=("3", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "1": ["1"], + "2": [ + ["2|Control(send=Send(node='2', arg=3))"], + ["2|Control(send=Send(node='flaky', arg=4))"], + ["2|3"], + ], + "flaky": ["flaky|4"], + }, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + ), + StateSnapshot( + values=["0"], + next=("1", "2", "2", "2", "flaky"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Control(send=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Control(send=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(Interrupt(value="Bahh", when="during"),), + state=None, + result=["flaky|4"], + ), + ), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": ["0"]}, + "thread_id": "1", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + ), + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_react_interrupt(checkpointer_name: str) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage @@ -2232,23 +2551,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 1, + "step": 0, "source": "loop", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, + "writes": None, "parents": {}, "thread_id": "2", }, @@ -2261,10 +2566,34 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), error=None, interrupts=(), state=None, @@ -2298,7 +2627,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 2, + "step": 1, "source": "update", "writes": { "agent": { @@ -2381,23 +2710,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 1, + "step": 0, "source": "loop", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, + "writes": None, "parents": {}, "thread_id": "3", }, @@ -2410,10 +2725,32 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), error=None, interrupts=(), state=None, @@ -2468,7 +2805,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 2, + "step": 1, "source": "update", "writes": { "agent": { @@ -2500,7 +2837,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", (), 0, AnyStr()), error=None, interrupts=(), state=None, @@ -2676,23 +3013,9 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: } }, metadata={ - "step": 1, + "step": 0, "source": "loop", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, + "writes": None, "parents": {}, "thread_id": "2", }, @@ -2705,10 +3028,34 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: } }, tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", 0), + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), error=None, interrupts=(), state=None, @@ -2742,7 +3089,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: } }, metadata={ - "step": 2, + "step": 1, "source": "update", "writes": { "agent": { @@ -2833,7 +3180,7 @@ async def test_max_concurrency(checkpointer_name: str) -> None: graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"]) thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1) == ["0", "1"] + assert await graph.ainvoke(["0"], thread1, debug=True) == ["0", "1"] state = await graph.aget_state(thread1) assert state.values == ["0", "1"] assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] @@ -5477,7 +5824,33 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=( @@ -5486,22 +5859,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, + "step": 0, + "writes": None, "thread_id": "1", }, parent_config=[ @@ -5533,14 +5892,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 2, + "step": 1, "writes": { "agent": { "messages": AIMessage( @@ -5636,8 +5995,40 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ] }, tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), ), next=("tools", "tools"), config=tup.config, @@ -5645,25 +6036,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 4, + "step": 2, "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), }, }, "thread_id": "1", @@ -5710,7 +6090,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "update", - "step": 5, + "step": 3, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -5773,7 +6153,35 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=( @@ -5782,22 +6190,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, + "step": 0, + "writes": None, "thread_id": "2", }, parent_config=[ @@ -5829,14 +6223,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 2, + "step": 1, "writes": { "agent": { "messages": AIMessage( @@ -5932,8 +6326,42 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ] }, tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), ), next=("tools", "tools"), config=tup.config, @@ -5941,25 +6369,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 4, + "step": 2, "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), }, }, "thread_id": "2", @@ -6006,7 +6423,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "update", - "step": 5, + "step": 3, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -8773,7 +9190,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert await app.ainvoke({"my_key": ""}, config, debug=True) == { - "my_key": "", + "my_key": " and parallel", } assert await app.ainvoke(None, config, debug=True) == { @@ -8808,6 +9225,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: c async for c in app.astream({"my_key": ""}, config, stream_mode="values") ] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -8826,6 +9244,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # while we're waiting for the node w/ interrupt inside to finish assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -8840,6 +9259,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: c async for c in app.astream({"my_key": ""}, config, stream_mode="values") ] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -10580,7 +11000,8 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: # invoke and pause at nested interrupt assert await graph.ainvoke( - {"subjects": ["cats", "dogs"]}, config={**config, "callbacks": [tracer]} + {"subjects": ["cats", "dogs"]}, + config={**config, "callbacks": [tracer]}, ) == { "subjects": ["cats", "dogs"], "jokes": [], @@ -10592,10 +11013,19 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: assert outer_state == StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), PregelTask( AnyStr(), "generate_joke", - (PUSH, 0), + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -10606,7 +11036,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, 1), + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -10625,24 +11055,25 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: }, metadata={ "parents": {}, - "source": "loop", - "writes": None, - "step": 0, + "source": "input", + "writes": { + "__start__": { + "subjects": [ + "cats", + "dogs", + ], + } + }, + "step": -1, "thread_id": "1", }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, + parent_config=None, ) # update state of dogs joke graph await graph.aupdate_state( - outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + outer_state.tasks[2].state, {"subject": "turtles - hohoho"} ) # continue past interrupt @@ -10675,7 +11106,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, "thread_id": "1", }, created_at=AnyStr(), @@ -10718,7 +11149,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, "thread_id": "1", }, created_at=AnyStr(), @@ -10731,13 +11162,22 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: }, ), StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - next=("generate_joke", "generate_joke"), + values={"jokes": []}, + next=("__start__", "generate_joke", "generate_joke"), tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), PregelTask( AnyStr(), "generate_joke", - (PUSH, 0), + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -10749,7 +11189,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, 1), + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -10766,40 +11206,6 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: "checkpoint_id": AnyStr(), } }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"subjects": ["cats", "dogs"]}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, metadata={ "parents": {}, "source": "input", @@ -10811,7 +11217,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: parent_config=None, ), ] - assert actual_history[1] == expected_history[1] + assert actual_history == expected_history @pytest.mark.skipif( diff --git a/libs/scheduler-kafka/Makefile b/libs/scheduler-kafka/Makefile index 5d899a3ea..8d62c9df2 100644 --- a/libs/scheduler-kafka/Makefile +++ b/libs/scheduler-kafka/Makefile @@ -19,7 +19,7 @@ test: exit $$EXIT_CODE test_watch: - make start-services && poetry run ptw . -- $(TEST_PATH); \ + make start-services && poetry run ptw . -- -x $(TEST_PATH); \ EXIT_CODE=$$?; \ make stop-services; \ exit $$EXIT_CODE diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index f4ffc8820..970d55be8 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -38,7 +38,7 @@ from langgraph.scheduler.kafka.types import ( Sendable, Topics, ) -from langgraph.types import LoopProtocol, RetryPolicy +from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable @@ -198,6 +198,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): msg["task"]["path"], msg["task"]["id"], checkpoint=saved.checkpoint, + pending_writes=saved.pending_writes or [], processes=graph.nodes, channels=channels, managed=managed, @@ -211,6 +212,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): runner = PregelRunner( submit=submit, put_writes=partial(self._put_writes, submit, msg["config"]), + schedule_task=self._schedule_task, ) async for _ in runner.atick([task], reraise=False): pass @@ -239,6 +241,14 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): ) await fut + def _schedule_task( + self, + task: PregelExecutableTask, + idx: int, + ) -> None: + # will be scheduled by orchestrator when executor finishes + pass + def _put_writes( self, submit: Submit, @@ -400,6 +410,7 @@ class KafkaExecutor(AbstractContextManager): msg["task"]["path"], msg["task"]["id"], checkpoint=saved.checkpoint, + pending_writes=saved.pending_writes or [], processes=graph.nodes, channels=channels, managed=managed, @@ -412,6 +423,7 @@ class KafkaExecutor(AbstractContextManager): runner = PregelRunner( submit=submit, put_writes=partial(self._put_writes, submit, msg["config"]), + schedule_task=self._schedule_task, ) for _ in runner.tick([task], reraise=False): pass @@ -440,6 +452,14 @@ class KafkaExecutor(AbstractContextManager): ) fut.result() + def _schedule_task( + self, + task: PregelExecutableTask, + idx: int, + ) -> None: + # will be scheduled by orchestrator when executor finishes + pass + def _put_writes( self, submit: Submit, diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 493b02d42..4e5be8470 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -161,18 +161,18 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): specs=graph.channels, output_keys=graph.output_channels, stream_keys=graph.stream_channels, + interrupt_after=graph.interrupt_after_nodes, + interrupt_before=graph.interrupt_before_nodes, check_subgraphs=False, ) as loop: - if loop.tick( - input_keys=graph.input_channels, - interrupt_after=graph.interrupt_after_nodes, - interrupt_before=graph.interrupt_before_nodes, - ): + if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved if hasattr(loop, "_put_checkpoint_fut"): await loop._put_checkpoint_fut # schedule any new tasks - if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]: + if new_tasks := [ + t for t in loop.tasks.values() if not t.scheduled and not t.writes + ]: # send messages to executor futures = await asyncio.gather( *( @@ -351,18 +351,18 @@ class KafkaOrchestrator(AbstractContextManager): specs=graph.channels, output_keys=graph.output_channels, stream_keys=graph.stream_channels, + interrupt_after=graph.interrupt_after_nodes, + interrupt_before=graph.interrupt_before_nodes, check_subgraphs=False, ) as loop: - if loop.tick( - input_keys=graph.input_channels, - interrupt_after=graph.interrupt_after_nodes, - interrupt_before=graph.interrupt_before_nodes, - ): + if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved if hasattr(loop, "_put_checkpoint_fut"): loop._put_checkpoint_fut.result() # schedule any new tasks - if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]: + if new_tasks := [ + t for t in loop.tasks.values() if not t.scheduled and not t.writes + ]: # send messages to executor futures = [ self.producer.send( diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py index 8230960b4..8a109631b 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py @@ -24,8 +24,8 @@ class MessageToOrchestrator(TypedDict): class ExecutorTask(TypedDict): - id: str - path: tuple[str, ...] + id: Optional[str] + path: tuple[Union[str, int], ...] class MessageToExecutor(TypedDict): diff --git a/libs/scheduler-kafka/tests/test_push.py b/libs/scheduler-kafka/tests/test_push.py new file mode 100644 index 000000000..7d7f48e95 --- /dev/null +++ b/libs/scheduler-kafka/tests/test_push.py @@ -0,0 +1,206 @@ +import operator +from typing import ( + Annotated, + Literal, + Union, +) + +import pytest +from aiokafka import AIOKafkaProducer + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.constants import START +from langgraph.errors import NodeInterrupt +from langgraph.graph.state import CompiledStateGraph, StateGraph +from langgraph.scheduler.kafka import serde +from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics +from langgraph.types import Control, Send +from tests.any import AnyDict +from tests.drain import drain_topics_async + +pytestmark = pytest.mark.anyio + + +def mk_push_graph( + checkpointer: BaseCheckpointSaver, +) -> CompiledStateGraph: + # copied from test_send_dedupe_on_resume + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + self.__name__ = name + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Control): + state.state = update + return state + else: + return update + + def send_for_fun(state): + return [ + Send("2", Control(send=Send("2", 3))), + Send("2", Control(send=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + return builder.compile(checkpointer=checkpointer) + + +async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: + input = ["0"] + config = {"configurable": {"thread_id": "1"}} + graph = mk_push_graph(acheckpointer) + graph_compare = mk_push_graph(acheckpointer) + + # start a new run + async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: + await producer.send_and_wait( + topics.orchestrator, + MessageToOrchestrator(input=input, config=config), + ) + + # drain topics + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) + + # check state + state = await graph.aget_state(config) + assert all(not t.error for t in state.tasks) + assert state.next == ("flaky",) + assert ( + state.values + == await graph_compare.ainvoke(input, {"configurable": {"thread_id": "2"}}) + == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + ] + ) + + # check history + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 2 + + # check messages + assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_send": None, + } + for c in reversed(history) + for _ in c.tasks + ] + assert exec_msgs == [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": _convert_path(t.path), + }, + "finally_send": None, + } + for c in reversed(history) + for t in c.tasks + ] + + # resume the thread + async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: + await producer.send_and_wait( + topics.orchestrator, + MessageToOrchestrator(input=None, config=config), + ) + + orch_msgs, exec_msgs = await drain_topics_async(topics, graph) + + # check final state + state = await graph.aget_state(config) + assert state.next == () + assert ( + state.values + == await graph_compare.ainvoke(None, {"configurable": {"thread_id": "2"}}) + == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + ) + + # check history + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 4 + + # check executions + # node "2" doesn't get called again, as we recover writes saved before + assert graph.builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert graph.builder.nodes["flaky"].runnable.func.ticks == 2 + + +def _convert_path( + path: tuple[Union[str, int, tuple], ...], +) -> list[Union[str, int, list]]: + return list(_convert_path(p) if isinstance(p, tuple) else p for p in path) diff --git a/libs/scheduler-kafka/tests/test_push_sync.py b/libs/scheduler-kafka/tests/test_push_sync.py new file mode 100644 index 000000000..51e60d51c --- /dev/null +++ b/libs/scheduler-kafka/tests/test_push_sync.py @@ -0,0 +1,208 @@ +import operator +from typing import ( + Annotated, + Literal, + Union, +) + +import pytest + +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.constants import START +from langgraph.errors import NodeInterrupt +from langgraph.graph.state import CompiledStateGraph, StateGraph +from langgraph.scheduler.kafka import serde +from langgraph.scheduler.kafka.default_sync import DefaultProducer +from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics +from langgraph.types import Control, Send +from tests.any import AnyDict +from tests.drain import drain_topics + +pytestmark = pytest.mark.anyio + + +def mk_push_graph( + checkpointer: BaseCheckpointSaver, +) -> CompiledStateGraph: + # copied from test_send_dedupe_on_resume + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + self.__name__ = name + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Control): + state.state = update + return state + else: + return update + + def send_for_fun(state): + return [ + Send("2", Control(send=Send("2", 3))), + Send("2", Control(send=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + return builder.compile(checkpointer=checkpointer) + + +def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: + input = ["0"] + config = {"configurable": {"thread_id": "1"}} + graph = mk_push_graph(acheckpointer) + graph_compare = mk_push_graph(acheckpointer) + + # start a new run + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=input, config=config)), + ) + producer.flush() + + # drain topics + orch_msgs, exec_msgs = drain_topics(topics, graph) + + # check state + state = graph.get_state(config) + assert all(not t.error for t in state.tasks) + assert state.next == ("flaky",) + assert ( + state.values + == graph_compare.invoke(input, {"configurable": {"thread_id": "2"}}) + == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + ] + ) + + # check history + history = [c for c in graph.get_state_history(config)] + assert len(history) == 2 + + # check messages + assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_send": None, + } + for c in reversed(history) + for _ in c.tasks + ] + assert exec_msgs == [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": _convert_path(t.path), + }, + "finally_send": None, + } + for c in reversed(history) + for t in c.tasks + ] + + # resume the thread + with DefaultProducer() as producer: + producer.send( + topics.orchestrator, + value=serde.dumps(MessageToOrchestrator(input=None, config=config)), + ) + producer.flush() + + orch_msgs, exec_msgs = drain_topics(topics, graph) + + # check final state + state = graph.get_state(config) + assert state.next == () + assert ( + state.values + == graph_compare.invoke(None, {"configurable": {"thread_id": "2"}}) + == [ + "0", + "1", + "2|Control(send=Send(node='2', arg=3))", + "2|Control(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + ) + + # check history + history = [c for c in graph.get_state_history(config)] + assert len(history) == 4 + + # check executions + # node "2" doesn't get called again, as we recover writes saved before + assert graph.builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert graph.builder.nodes["flaky"].runnable.func.ticks == 2 + + +def _convert_path( + path: tuple[Union[str, int, tuple], ...], +) -> list[Union[str, int, list]]: + return list(_convert_path(p) if isinstance(p, tuple) else p for p in path) From 3ad966e05737a840564d849e72178d97d7ef118e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 14:08:59 -0800 Subject: [PATCH 02/22] Update --- libs/langgraph/tests/conftest.py | 1 - libs/langgraph/tests/test_pregel.py | 35 ++++++------ libs/langgraph/tests/test_pregel_async.py | 69 +++++++++++++++-------- 3 files changed, 61 insertions(+), 44 deletions(-) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 6bc23f907..f22f96de5 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -335,7 +335,6 @@ ALL_CHECKPOINTERS_SYNC = [ ALL_CHECKPOINTERS_ASYNC = [ "memory", "sqlite_aio", - "duckdb_aio", "postgres_aio", "postgres_aio_pipe", "postgres_aio_pool", diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bf92f489c..1d0c0b52e 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1875,16 +1875,15 @@ def test_send_dedupe_on_resume( if isinstance(state, list) else ["|".join((self.name, str(state)))] ) - if isinstance(state, Control): - state.state = update - return state + if isinstance(state, GraphCommand): + return state.copy(update=update) else: return update def send_for_fun(state): return [ - Send("2", Control(send=Send("2", 3))), - Send("2", Control(send=Send("flaky", 4))), + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("flaky", 4))), "3.1", ] @@ -1906,8 +1905,8 @@ def test_send_dedupe_on_resume( assert graph.invoke(["0"], thread1, debug=1) == [ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", ] assert builder.nodes["2"].runnable.func.ticks == 3 @@ -1922,8 +1921,8 @@ def test_send_dedupe_on_resume( assert graph.invoke(None, thread1, debug=1) == [ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", "flaky|4", "3", @@ -1945,8 +1944,8 @@ def test_send_dedupe_on_resume( values=[ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", "flaky|4", "3", @@ -1981,8 +1980,8 @@ def test_send_dedupe_on_resume( values=[ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", "flaky|4", ], @@ -1999,8 +1998,8 @@ def test_send_dedupe_on_resume( "writes": { "1": ["1"], "2": [ - ["2|Control(send=Send(node='2', arg=3))"], - ["2|Control(send=Send(node='flaky', arg=4))"], + ["2|Command(send=Send(node='2', arg=3))"], + ["2|Command(send=Send(node='flaky', arg=4))"], ["2|3"], ], "flaky": ["flaky|4"], @@ -2085,7 +2084,7 @@ def test_send_dedupe_on_resume( error=None, interrupts=(), state=None, - result=["2|Control(send=Send(node='2', arg=3))"], + result=["2|Command(send=Send(node='2', arg=3))"], ), PregelTask( id=AnyStr(), @@ -2099,7 +2098,7 @@ def test_send_dedupe_on_resume( error=None, interrupts=(), state=None, - result=["2|Control(send=Send(node='flaky', arg=4))"], + result=["2|Command(send=Send(node='flaky', arg=4))"], ), PregelTask( id=AnyStr(), @@ -2904,7 +2903,7 @@ def test_send_react_interrupt_control( # interrupt-update-resume flow, creating new Send in update call - # TODO add here test with invoke(Control()) + # TODO add here test with invoke(Command()) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b32100a2e..cc1bd0d39 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2035,7 +2035,9 @@ async def test_concurrent_emit_sends() -> None: ] -async def test_send_sequences() -> None: +@pytest.mark.repeat(10) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_sequences(checkpointer_name: str) -> None: class Node: def __init__(self, name: str): self.name = name @@ -2074,22 +2076,40 @@ async def test_send_sequences() -> None: assert await graph.ainvoke(["0"]) == [ "0", "1", - "3.1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='2', arg=4))", - "3", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", "2|3", "2|4", "3", + "3.1", ] + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + ] + assert await graph.ainvoke(None, thread1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + "3", + "3.1", + ] + @pytest.mark.repeat(20) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: - if checkpointer_name == "duckdb_aio": - pytest.skip("DuckDB isn't returning the right history") - class InterruptOnce: ticks: int = 0 @@ -2112,16 +2132,15 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: if isinstance(state, list) else ["|".join((self.name, str(state)))] ) - if isinstance(state, Control): - state.state = update - return state + if isinstance(state, GraphCommand): + return state.copy(update=update) else: return update def send_for_fun(state): return [ - Send("2", Control(send=Send("2", 3))), - Send("2", Control(send=Send("flaky", 4))), + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("flaky", 4))), "3.1", ] @@ -2144,8 +2163,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert await graph.ainvoke(["0"], thread1, debug=1) == [ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", ] assert builder.nodes["2"].runnable.func.ticks == 3 @@ -2154,8 +2173,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert await graph.ainvoke(None, thread1, debug=1) == [ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", "flaky|4", "3", @@ -2172,8 +2191,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: values=[ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", "flaky|4", "3", @@ -2208,8 +2227,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: values=[ "0", "1", - "2|Control(send=Send(node='2', arg=3))", - "2|Control(send=Send(node='flaky', arg=4))", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", "2|3", "flaky|4", ], @@ -2226,8 +2245,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: "writes": { "1": ["1"], "2": [ - ["2|Control(send=Send(node='2', arg=3))"], - ["2|Control(send=Send(node='flaky', arg=4))"], + ["2|Command(send=Send(node='2', arg=3))"], + ["2|Command(send=Send(node='flaky', arg=4))"], ["2|3"], ], "flaky": ["flaky|4"], @@ -2312,7 +2331,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: error=None, interrupts=(), state=None, - result=["2|Control(send=Send(node='2', arg=3))"], + result=["2|Command(send=Send(node='2', arg=3))"], ), PregelTask( id=AnyStr(), @@ -2326,7 +2345,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: error=None, interrupts=(), state=None, - result=["2|Control(send=Send(node='flaky', arg=4))"], + result=["2|Command(send=Send(node='flaky', arg=4))"], ), PregelTask( id=AnyStr(), From 0e872e74828e7d6ba328310e41b50fd509aec63d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 14:12:33 -0800 Subject: [PATCH 03/22] Lin t --- libs/langgraph/tests/test_pregel.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 1d0c0b52e..c57948c00 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1834,13 +1834,12 @@ def test_send_sequences() -> None: assert graph.invoke(["0"]) == [ "0", "1", - "3.1", "2|Command(send=Send(node='2', arg=3))", "2|Command(send=Send(node='2', arg=4))", - "3", "2|3", "2|4", "3", + "3.1", ] From 090b53ccc1970adb114fc1541da6a2577b1ffdf8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 14:14:23 -0800 Subject: [PATCH 04/22] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 76020cf89..a965342f2 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -887,6 +887,7 @@ class Pregel(PregelProtocol): # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, + saved.pending_writes or [], self.nodes, channels, managed, @@ -1106,6 +1107,7 @@ class Pregel(PregelProtocol): # tasks for this checkpoint next_tasks = prepare_next_tasks( checkpoint, + saved.pending_writes or [], self.nodes, channels, managed, From ea64ac5c07208be50bbb428fa5b5c792b70c6a39 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 14:18:14 -0800 Subject: [PATCH 05/22] Update --- libs/scheduler-kafka/tests/test_push.py | 13 ++++++------- libs/scheduler-kafka/tests/test_push_sync.py | 13 ++++++------- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/libs/scheduler-kafka/tests/test_push.py b/libs/scheduler-kafka/tests/test_push.py index 7d7f48e95..d7611cd00 100644 --- a/libs/scheduler-kafka/tests/test_push.py +++ b/libs/scheduler-kafka/tests/test_push.py @@ -11,10 +11,10 @@ from aiokafka import AIOKafkaProducer from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import START from langgraph.errors import NodeInterrupt -from langgraph.graph.state import CompiledStateGraph, StateGraph +from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from langgraph.types import Control, Send +from langgraph.types import Send from tests.any import AnyDict from tests.drain import drain_topics_async @@ -48,16 +48,15 @@ def mk_push_graph( if isinstance(state, list) else ["|".join((self.name, str(state)))] ) - if isinstance(state, Control): - state.state = update - return state + if isinstance(state, GraphCommand): + return state.copy(update=update) else: return update def send_for_fun(state): return [ - Send("2", Control(send=Send("2", 3))), - Send("2", Control(send=Send("flaky", 4))), + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("flaky", 4))), "3.1", ] diff --git a/libs/scheduler-kafka/tests/test_push_sync.py b/libs/scheduler-kafka/tests/test_push_sync.py index 51e60d51c..09c84c782 100644 --- a/libs/scheduler-kafka/tests/test_push_sync.py +++ b/libs/scheduler-kafka/tests/test_push_sync.py @@ -10,11 +10,11 @@ import pytest from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import START from langgraph.errors import NodeInterrupt -from langgraph.graph.state import CompiledStateGraph, StateGraph +from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.default_sync import DefaultProducer from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from langgraph.types import Control, Send +from langgraph.types import Send from tests.any import AnyDict from tests.drain import drain_topics @@ -48,16 +48,15 @@ def mk_push_graph( if isinstance(state, list) else ["|".join((self.name, str(state)))] ) - if isinstance(state, Control): - state.state = update - return state + if isinstance(state, GraphCommand): + return state.copy(update=update) else: return update def send_for_fun(state): return [ - Send("2", Control(send=Send("2", 3))), - Send("2", Control(send=Send("flaky", 4))), + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("flaky", 4))), "3.1", ] From d0567dc7be19ba102a6b161829c2bd75b7e6f7d4 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 15:38:48 -0800 Subject: [PATCH 06/22] Add feature flag (default off) so we can merge this before releasing - Add additional ci job to test with FF on --- .github/workflows/_test_langgraph.yml | 6 ++ libs/langgraph/Makefile | 2 +- libs/langgraph/langgraph/constants.py | 3 + libs/langgraph/langgraph/pregel/algo.py | 2 +- libs/langgraph/langgraph/pregel/write.py | 8 +- libs/langgraph/tests/conftest.py | 1 - libs/langgraph/tests/test_pregel.py | 112 +++++++++++++++++----- libs/langgraph/tests/test_pregel_async.py | 109 ++++++++++++++++----- 8 files changed, 191 insertions(+), 52 deletions(-) diff --git a/.github/workflows/_test_langgraph.yml b/.github/workflows/_test_langgraph.yml index b8e2679fe..b771da867 100644 --- a/.github/workflows/_test_langgraph.yml +++ b/.github/workflows/_test_langgraph.yml @@ -19,9 +19,13 @@ jobs: - "3.13" core-version: - "latest" + ff-send-v2: + - "false" include: - python-version: "3.11" core-version: ">=0.2.42,<0.3.0" + - python-version: "3.11" + ff-send-v2: "true" defaults: run: @@ -52,6 +56,8 @@ jobs: - name: Run tests shell: bash + env: + LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }} run: | make test diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 1e249a0cd..2aacf6db8 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -49,7 +49,7 @@ test: exit $$EXIT_CODE test_watch: - make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ + make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 5d18b5262..c987dd69a 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,4 +1,5 @@ import sys +from os import getenv from types import MappingProxyType from typing import Any, Literal, Mapping, cast @@ -81,6 +82,8 @@ NS_END = sys.intern(":") # for checkpoint_ns, for each level, separates the namespace from the task_id CONF = cast(Literal["configurable"], sys.intern("configurable")) # key for the configurable dict in RunnableConfig +FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true" +# temporary flag to enable new Send semantics RESERVED = { TAG_HIDDEN, diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 6ae0253ac..d6669bdb5 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -173,7 +173,7 @@ def local_write( """Function injected under CONFIG_KEY_SEND in task config, to write to channels. Validates writes and forwards them to `commit` function.""" for chan, value in writes: - if chan == PUSH: + if chan in (PUSH, TASKS): if not isinstance(value, Send): raise InvalidUpdateError(f"Expected Send, got {value}") if value.node not in process_keys: diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index ba783453c..3af0fe5e9 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -14,7 +14,7 @@ from typing import ( from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec -from langgraph.constants import CONF, CONFIG_KEY_SEND, PUSH, TASKS, Send +from langgraph.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send from langgraph.errors import InvalidUpdateError from langgraph.utils.runnable import RunnableCallable @@ -119,7 +119,11 @@ class ChannelWrite(RunnableCallable): if w.value is PASSTHROUGH: raise InvalidUpdateError("PASSTHROUGH value must be replaced") # split packets and entries - sends = [(PUSH, packet) for packet in writes if isinstance(packet, Send)] + sends = [ + (PUSH if FF_SEND_V2 else TASKS, packet) + for packet in writes + if isinstance(packet, Send) + ] entries = [write for write in writes if isinstance(write, ChannelWriteEntry)] # process entries into values values = [ diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index f22f96de5..eae7694ff 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -327,7 +327,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: ALL_CHECKPOINTERS_SYNC = [ "memory", "sqlite", - "duckdb", "postgres", "postgres_pipe", "postgres_pool", diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c57948c00..bce736ac5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -54,7 +54,14 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START +from langgraph.constants import ( + CONFIG_KEY_NODE_FINISHED, + ERROR, + FF_SEND_V2, + PULL, + PUSH, + START, +) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt from langgraph.graph import END, Graph, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages @@ -1781,17 +1788,31 @@ def test_concurrent_emit_sends() -> None: builder.add_conditional_edges("1.1", send_for_profit) builder.add_conditional_edges("2", route_to_three) graph = builder.compile() - assert graph.invoke(["0"]) == [ - "0", - "1", - "1.1", - "2|1", - "2|2", - "2|3", - "2|4", - "3", - "3.1", - ] + assert graph.invoke(["0"]) == ( + [ + "0", + "1", + "1.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] + ) def test_send_sequences() -> None: @@ -1831,16 +1852,31 @@ def test_send_sequences() -> None: builder.add_conditional_edges("1", send_for_fun) builder.add_conditional_edges("2", route_to_three) graph = builder.compile() - assert graph.invoke(["0"]) == [ - "0", - "1", - "2|Command(send=Send(node='2', arg=3))", - "2|Command(send=Send(node='2', arg=4))", - "2|3", - "2|4", - "3", - "3.1", - ] + assert ( + graph.invoke(["0"]) + == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "3.1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + ) @pytest.mark.repeat(20) @@ -1848,8 +1884,8 @@ def test_send_sequences() -> None: def test_send_dedupe_on_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - if checkpointer_name == "duckdb": - pytest.skip("DuckDB isn't returning the right history") + if not FF_SEND_V2: + pytest.skip("Send deduplication is only available in Send V2") checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class InterruptOnce: @@ -2298,6 +2334,9 @@ def test_send_react_interrupt( } assert foo_called == 0 + if not FF_SEND_V2: + return + # get state should show the pending task state = graph.get_state(thread1) assert state == StateSnapshot( @@ -2741,6 +2780,9 @@ def test_send_react_interrupt_control( } assert foo_called == 1 + if not FF_SEND_V2: + return + # interrupt-update-resume flow foo_called = 0 graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) @@ -5866,6 +5908,9 @@ def test_state_graph_packets( {"__interrupt__": ()}, ] + if not FF_SEND_V2: + return + assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "messages": [ @@ -12264,6 +12309,21 @@ def test_send_to_nested_graphs( # check state outer_state = graph.get_state(config) + + if not FF_SEND_V2: + # update state of dogs joke graph + graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) + + # continue past interrupt + assert sorted( + graph.stream(None, config=config), + key=lambda d: d["generate_joke"]["jokes"][0], + ) == [ + {"generate_joke": {"jokes": ["Joke about cats - hohoho"]}}, + {"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}}, + ] + return + assert outer_state == StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, tasks=( @@ -12409,7 +12469,9 @@ def test_send_to_nested_graphs( tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), ) # update state of dogs joke graph - graph.update_state(outer_state.tasks[2].state, {"subject": "turtles - hohoho"}) + graph.update_state( + outer_state.tasks[2 if FF_SEND_V2 else 1].state, {"subject": "turtles - hohoho"} + ) # continue past interrupt assert sorted( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index cc1bd0d39..c82935082 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -51,7 +51,14 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START +from langgraph.constants import ( + CONFIG_KEY_NODE_FINISHED, + ERROR, + FF_SEND_V2, + PULL, + PUSH, + START, +) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt from langgraph.graph import END, Graph, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages @@ -2022,17 +2029,31 @@ async def test_concurrent_emit_sends() -> None: builder.add_conditional_edges("1.1", send_for_profit) builder.add_conditional_edges("2", route_to_three) graph = builder.compile() - assert await graph.ainvoke(["0"]) == [ - "0", - "1", - "1.1", - "2|1", - "2|2", - "2|3", - "2|4", - "3", - "3.1", - ] + assert await graph.ainvoke(["0"]) == ( + [ + "0", + "1", + "1.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] + ) @pytest.mark.repeat(10) @@ -2073,16 +2094,34 @@ async def test_send_sequences(checkpointer_name: str) -> None: builder.add_conditional_edges("1", send_for_fun) builder.add_conditional_edges("2", route_to_three) graph = builder.compile() - assert await graph.ainvoke(["0"]) == [ - "0", - "1", - "2|Command(send=Send(node='2', arg=3))", - "2|Command(send=Send(node='2', arg=4))", - "2|3", - "2|4", - "3", - "3.1", - ] + assert ( + await graph.ainvoke(["0"]) + == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "3.1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + ) + + if not FF_SEND_V2: + return async with awith_checkpointer(checkpointer_name) as checkpointer: graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"]) @@ -2110,6 +2149,9 @@ async def test_send_sequences(checkpointer_name: str) -> None: @pytest.mark.repeat(20) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: + if not FF_SEND_V2: + pytest.skip("Send deduplication is only available in Send V2") + class InterruptOnce: ticks: int = 0 @@ -2542,6 +2584,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } assert foo_called == 0 + if not FF_SEND_V2: + return + # get state should show the pending task state = await graph.aget_state(thread1) assert state == StateSnapshot( @@ -3004,6 +3049,9 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: } assert foo_called == 0 + if not FF_SEND_V2: + return + # get state should show the pending task state = await graph.aget_state(thread1) assert state == StateSnapshot( @@ -5826,6 +5874,9 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: {"__interrupt__": ()}, ] + if not FF_SEND_V2: + return + assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ "messages": [ @@ -11029,6 +11080,20 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: # check state outer_state = await graph.aget_state(config) + + if not FF_SEND_V2: + # update state of dogs joke graph + await graph.aupdate_state( + outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + ) + + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], + } + return + assert outer_state == StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, tasks=( From 2ff49d22002da9d0b09136d09ef3e5b22889e816 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 15:43:20 -0800 Subject: [PATCH 07/22] Update --- .github/workflows/_test_langgraph.yml | 3 ++- libs/scheduler-kafka/tests/test_push.py | 5 ++++- libs/scheduler-kafka/tests/test_push_sync.py | 5 ++++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/.github/workflows/_test_langgraph.yml b/.github/workflows/_test_langgraph.yml index b771da867..5c3f5182e 100644 --- a/.github/workflows/_test_langgraph.yml +++ b/.github/workflows/_test_langgraph.yml @@ -25,12 +25,13 @@ jobs: - python-version: "3.11" core-version: ">=0.2.42,<0.3.0" - python-version: "3.11" + core-version: "latest" ff-send-v2: "true" defaults: run: working-directory: libs/langgraph - name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }})" + name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }}, ff-send-v2: ${{ matrix.ff-send-v2 }})" steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }} diff --git a/libs/scheduler-kafka/tests/test_push.py b/libs/scheduler-kafka/tests/test_push.py index d7611cd00..15e9211a2 100644 --- a/libs/scheduler-kafka/tests/test_push.py +++ b/libs/scheduler-kafka/tests/test_push.py @@ -9,7 +9,7 @@ import pytest from aiokafka import AIOKafkaProducer from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import START +from langgraph.constants import FF_SEND_V2, START from langgraph.errors import NodeInterrupt from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph from langgraph.scheduler.kafka import serde @@ -77,6 +77,9 @@ def mk_push_graph( async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: + if not FF_SEND_V2: + pytest.skip("Test requires FF_SEND_V2") + input = ["0"] config = {"configurable": {"thread_id": "1"}} graph = mk_push_graph(acheckpointer) diff --git a/libs/scheduler-kafka/tests/test_push_sync.py b/libs/scheduler-kafka/tests/test_push_sync.py index 09c84c782..27cd96cb7 100644 --- a/libs/scheduler-kafka/tests/test_push_sync.py +++ b/libs/scheduler-kafka/tests/test_push_sync.py @@ -8,7 +8,7 @@ from typing import ( import pytest from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import START +from langgraph.constants import FF_SEND_V2, START from langgraph.errors import NodeInterrupt from langgraph.graph.state import CompiledStateGraph, GraphCommand, StateGraph from langgraph.scheduler.kafka import serde @@ -77,6 +77,9 @@ def mk_push_graph( def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: + if not FF_SEND_V2: + pytest.skip("Test requires FF_SEND_V2") + input = ["0"] config = {"configurable": {"thread_id": "1"}} graph = mk_push_graph(acheckpointer) From 810ae0ef513c731c428ca89129912792c9e93a2c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:44:03 -0800 Subject: [PATCH 08/22] lib: Add interrupt() function - This works similarly to the input() function from stdlib - calling it in a node interrupts execution - invoking the graph with Command(resume=...) will set ... as the return value of interrupt() so that the node can access the "answer" to the "question" - This PR also starts the work to control the graph on invoke/stream with Command() input, to be continued in a future PR --- libs/langgraph/langgraph/constants.py | 8 +++ libs/langgraph/langgraph/graph/state.py | 3 +- libs/langgraph/langgraph/pregel/__init__.py | 64 ++++++++++++++++++ libs/langgraph/langgraph/pregel/algo.py | 59 +++++++++++++---- libs/langgraph/langgraph/pregel/io.py | 63 +++++++++++++++++- libs/langgraph/langgraph/pregel/loop.py | 73 ++++++++++++++++----- libs/langgraph/langgraph/types.py | 17 ++++- libs/langgraph/langgraph/utils/config.py | 7 ++ libs/langgraph/tests/test_pregel_async.py | 56 ++++++++++++---- 9 files changed, 306 insertions(+), 44 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index c987dd69a..478f23d68 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -11,6 +11,7 @@ from langgraph.types import Interrupt, Send # noqa: F401 # --- Empty read-only containers --- EMPTY_MAP: Mapping[str, Any] = MappingProxyType({}) EMPTY_SEQ: tuple[str, ...] = tuple() +MISSING = object() # --- Public constants --- TAG_NOSTREAM = sys.intern("langsmith:nostream") @@ -29,6 +30,8 @@ INPUT = sys.intern("__input__") # for values passed as input to the graph INTERRUPT = sys.intern("__interrupt__") # for dynamic interrupts raised by nodes +RESUME = sys.intern("__resume__") +# for values passed to resume a node after an interrupt ERROR = sys.intern("__error__") # for errors raised by nodes NO_WRITES = sys.intern("__no_writes__") @@ -70,6 +73,8 @@ CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns") # holds the current checkpoint_ns, "" for root graph CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished") # callback to be called when a node is finished +CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value") +# holds the value that "answers" an interrupt() call # --- Other constants --- PUSH = sys.intern("__pregel_push") @@ -84,12 +89,15 @@ CONF = cast(Literal["configurable"], sys.intern("configurable")) # key for the configurable dict in RunnableConfig FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true" # temporary flag to enable new Send semantics +NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000") +# the task_id to use for writes that are not associated with a task RESERVED = { TAG_HIDDEN, # reserved write keys INPUT, INTERRUPT, + RESUME, ERROR, NO_WRITES, SCHEDULED, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 2cae9fc7e..d3b95d9e2 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -92,8 +92,9 @@ class GraphCommand(Command, Generic[N]): update: Optional[dict[str, Any]] = None, goto: Union[str, Sequence[str]] = (), send: Union[Send, Sequence[Send]] = (), + resume: Optional[Union[Any, dict[str, Any]]] = None, ) -> None: - super().__init__(update=update, send=send) + super().__init__(update=update, send=send, resume=resume) self.goto = goto diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a965342f2..48306f842 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -66,9 +66,11 @@ from langgraph.constants import ( CONFIG_KEY_STREAM_WRITER, CONFIG_KEY_TASK_ID, ERROR, + INPUT, INTERRUPT, NS_END, NS_SEP, + NULL_TASK_ID, PUSH, SCHEDULED, ) @@ -519,6 +521,15 @@ class Pregel(PregelProtocol): config, subgraphs=True ) # apply pending writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): @@ -622,6 +633,15 @@ class Pregel(PregelProtocol): config, subgraphs=True ) # apply pending writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) if apply_pending_writes and saved.pending_writes: for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): @@ -898,6 +918,18 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer or None, manager=None, ) + # apply null writes + if null_writes := [ + w[1:] + for w in saved.pending_writes or [] + if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: if k in (ERROR, INTERRUPT, SCHEDULED): @@ -943,6 +975,16 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer or None, manager=None, ) + # apply null writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) # apply writes for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): @@ -1118,6 +1160,18 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer or None, manager=None, ) + # apply null writes + if null_writes := [ + w[1:] + for w in saved.pending_writes or [] + if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) # apply writes from tasks that already ran for tid, k, v in saved.pending_writes or []: if k in (ERROR, INTERRUPT, SCHEDULED): @@ -1163,6 +1217,16 @@ class Pregel(PregelProtocol): checkpointer=self.checkpointer or None, manager=None, ) + # apply null writes + if null_writes := [ + w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID + ]: + apply_writes( + saved.checkpoint, + channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + None, + ) for tid, k, v in saved.pending_writes: if k in (ERROR, INTERRUPT, SCHEDULED): continue diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index d6669bdb5..79ddb8722 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -36,17 +36,21 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, + CONFIG_KEY_RESUME_VALUE, CONFIG_KEY_SEND, CONFIG_KEY_STORE, CONFIG_KEY_TASK_ID, EMPTY_SEQ, INTERRUPT, + MISSING, NO_WRITES, NS_END, NS_SEP, + NULL_TASK_ID, PULL, PUSH, RESERVED, + RESUME, TAG_HIDDEN, TASKS, Send, @@ -199,6 +203,9 @@ def apply_writes( # any path parts after the 3rd are ignored for sorting # (we use them for eg. task ids which aren't good for sorting) tasks = sorted(tasks, key=lambda t: t.path[:3]) + # if no task has triggers this is applying writes from the null task only + # so we don't do anything other than update the channels written to + bump_step = any(t.triggers for t in tasks) # update seen versions for task in tasks: @@ -230,7 +237,7 @@ def apply_writes( ) # clear pending sends - if checkpoint["pending_sends"]: + if checkpoint["pending_sends"] and bump_step: checkpoint["pending_sends"].clear() # Group writes by channel @@ -238,12 +245,10 @@ def apply_writes( pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list) for task in tasks: for chan, val in task.writes: - if chan == NO_WRITES: + if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT): pass elif chan == TASKS: # TODO: remove branch in 1.0 checkpoint["pending_sends"].append(val) - elif chan == PUSH: - pass elif chan in channels: pending_writes_by_channel[chan].append(val) else: @@ -267,13 +272,14 @@ def apply_writes( updated_channels.add(chan) # Channels that weren't updated in this step are notified of a new step - for chan in channels: - if chan not in updated_channels: - if channels[chan].update([]) and get_next_version is not None: - checkpoint["channel_versions"][chan] = get_next_version( - max_version, - channels[chan], - ) + if bump_step: + for chan in channels: + if chan not in updated_channels: + if channels[chan].update([]) and get_next_version is not None: + checkpoint["channel_versions"][chan] = get_next_version( + max_version, + channels[chan], + ) # Return managed values writes to be applied externally return pending_writes_by_managed @@ -582,6 +588,14 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, + CONFIG_KEY_RESUME_VALUE: next( + ( + v + for tid, c, v in pending_writes + if tid in (NULL_TASK_ID, task_id) and c == RESUME + ), + MISSING, + ), }, ), triggers, @@ -599,6 +613,19 @@ def prepare_single_task( if name not in processes: return proc = processes[name] + print( + "preparing task", + task_path, + pending_writes, + sorted( + (chan, read_channel(channels, chan, return_exception=True)) + for chan in proc.triggers + # if not isinstance( + # read_channel(channels, chan, return_exception=True), + # EmptyChannelError, + # ) + ), + ) version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) null_version = version_type() # type: ignore[misc] if null_version is None: @@ -639,6 +666,7 @@ def prepare_single_task( "langgraph_path": task_path, "langgraph_checkpoint_ns": task_checkpoint_ns, } + print("preparing task", task_id, task_path, pending_writes) if task_id_checksum is not None: assert task_id == task_id_checksum if for_execution: @@ -691,6 +719,15 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, + CONFIG_KEY_RESUME_VALUE: next( + ( + v + for tid, c, v in pending_writes + if tid in (NULL_TASK_ID, task_id) + and c == RESUME + ), + MISSING, + ), }, ), triggers, diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 2a1f629cb..58a593645 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -1,11 +1,32 @@ from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypeVar, Union +from uuid import UUID from langchain_core.runnables.utils import AddableDict from langgraph.channels.base import BaseChannel, EmptyChannelError -from langgraph.constants import EMPTY_SEQ, ERROR, INTERRUPT, TAG_HIDDEN +from langgraph.constants import ( + EMPTY_SEQ, + ERROR, + FF_SEND_V2, + INTERRUPT, + NULL_TASK_ID, + PUSH, + RESUME, + TAG_HIDDEN, + TASKS, +) from langgraph.pregel.log import logger -from langgraph.types import PregelExecutableTask +from langgraph.types import Command, PregelExecutableTask, Send + + +def is_task_id(task_id: str) -> bool: + """Check if a string is a valid task id.""" + try: + u = UUID(task_id) + print(u.version) + except ValueError: + return False + return True def read_channel( @@ -44,6 +65,44 @@ def read_channels( return values +def map_command( + cmd: Command, +) -> Iterator[tuple[str, str, Any]]: + """Map input chunk to a sequence of pending writes in the form (channel, value).""" + if cmd.send: + if isinstance(cmd.send, (tuple, list)) and all( + isinstance(x, Send) + or isinstance(x, (list, tuple)) + and len(x) == 2 + and isinstance(x[0], str) + for x in cmd.send + ): + sends = cmd.send + else: + sends = [cmd.send] + for send in sends: + if isinstance(send, tuple) and len(send) == 2 and isinstance(send[0], str): + send = Send(*send) + if not isinstance(send, Send): + raise TypeError( + f"In Command.send, expected Send, got {type(send).__name__}" + ) + yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send) + if cmd.resume: + if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume): + for tid, resume in cmd.resume.items(): + yield (tid, RESUME, resume) + else: + yield (NULL_TASK_ID, RESUME, cmd.resume) + if cmd.update: + if not isinstance(cmd.update, dict): + raise TypeError( + f"Expected cmd.update to be a dict mapping channel names to update values, got {type(cmd.update).__name__}" + ) + for k, v in cmd.update.items(): + yield (NULL_TASK_ID, k, v) + + def map_input( input_channels: Union[str, Sequence[str]], chunk: Optional[Union[dict[str, Any], Any]], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 2e10a7a84..36811d02a 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,6 +1,6 @@ import asyncio import concurrent.futures -from collections import deque +from collections import defaultdict, deque from contextlib import AsyncExitStack, ExitStack from types import TracebackType from typing import ( @@ -52,7 +52,9 @@ from langgraph.constants import ( INPUT, INTERRUPT, NS_SEP, + NULL_TASK_ID, PUSH, + RESUME, SCHEDULED, TAG_HIDDEN, ) @@ -92,6 +94,7 @@ from langgraph.pregel.executor import ( Submit, ) from langgraph.pregel.io import ( + map_command, map_input, map_output_updates, map_output_values, @@ -102,7 +105,13 @@ from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.pregel.utils import get_new_channel_versions from langgraph.store.base import BaseStore -from langgraph.types import All, LoopProtocol, PregelExecutableTask, StreamProtocol +from langgraph.types import ( + All, + Command, + LoopProtocol, + PregelExecutableTask, + StreamProtocol, +) from langgraph.utils.config import patch_configurable V = TypeVar("V") @@ -273,7 +282,8 @@ class PregelLoop(LoopProtocol): task_id, ) # output writes - self._output_writes(task_id, writes) + if hasattr(self, "tasks"): + self._output_writes(task_id, writes) def accept_push( self, task: PregelExecutableTask, write_idx: int @@ -395,6 +405,19 @@ class PregelLoop(LoopProtocol): self.status = "out_of_steps" return False + # apply NULL writes + if null_writes := [ + w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID + ]: + mv_writes = apply_writes( + self.checkpoint, + self.channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + self.checkpointer_get_next_version, + ) + for key, values in mv_writes.items(): + self._update_mv(key, values) + print("applied null writes", null_writes) # prepare next tasks self.tasks = prepare_next_tasks( self.checkpoint, @@ -478,7 +501,7 @@ class PregelLoop(LoopProtocol): def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None: for tid, k, v in self.checkpoint_pending_writes: - if k in (ERROR, INTERRUPT): + if k in (ERROR, INTERRUPT, RESUME): continue if task := tasks.get(tid): if k == SCHEDULED: @@ -510,8 +533,21 @@ class PregelLoop(LoopProtocol): self._emit( "values", map_output_values, self.output_keys, True, self.channels ) + # map command to writes + elif isinstance(self.input, Command): + writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list) + # group writes by task ID + for tid, c, v in map_command(self.input): + writes[tid].append((c, v)) + if not writes: + raise EmptyInputError("Received empty Command input") + # save writes + for tid, ws in writes.items(): + self.put_writes(tid, ws) + print("applied cmd", writes) # map inputs to channel updates elif input_writes := deque(map_input(input_keys, self.input)): + # TODO shouldn't these writes be passed to put_writes too? # check if we should delegate (used by subgraphs in distributed mode) if self.config[CONF].get(CONFIG_KEY_DELEGATE): raise GraphDelegate( @@ -523,19 +559,22 @@ class PregelLoop(LoopProtocol): } ) # discard any unfinished tasks from previous checkpoint - discard_tasks = prepare_next_tasks( - self.checkpoint, - self.checkpoint_pending_writes, - self.nodes, - self.channels, - self.managed, - self.config, - self.step, - for_execution=True, - store=None, - checkpointer=None, - manager=None, - ) + if not isinstance(self.input, Command): + discard_tasks = prepare_next_tasks( + self.checkpoint, + self.checkpoint_pending_writes, + self.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + store=None, + checkpointer=None, + manager=None, + ) + else: + discard_tasks = {} # apply input writes mv_writes = apply_writes( self.checkpoint, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 23ea2d881..3337fc601 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -224,16 +224,18 @@ class Send: class Command: """One or more commands to update the graph's state and send messages to nodes.""" - __slots__ = ("update", "send") + __slots__ = ("update", "send", "resume") def __init__( self, *, update: Optional[dict[str, Any]] = None, send: Union[Send, Sequence[Send]] = (), + resume: Optional[Union[Any, dict[str, Any]]] = None, ) -> None: self.update = update self.send = send + self.resume = resume @property def __all_slots__(self) -> set[str]: @@ -307,3 +309,16 @@ class LoopProtocol: self.store = store self.step = step self.stop = stop + + +def interrupt(value: Any) -> Any: + from langgraph.constants import CONFIG_KEY_RESUME_VALUE, MISSING + from langgraph.errors import NodeInterrupt + from langgraph.utils.config import get_configurable + + conf = get_configurable() + print("interrupt", conf.get(CONFIG_KEY_RESUME_VALUE)) + if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING: + return resume + else: + raise NodeInterrupt(value) diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index fe25b6d9a..064a2b408 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -290,3 +290,10 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: ): empty["metadata"][key] = value return empty + + +def get_configurable() -> dict[str, Any]: + if var_config := var_child_runnable_config.get(): + return var_config[CONF] + else: + raise RuntimeError("Called get_configurable outside of a runnable context") diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c82935082..1e4823d7e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -69,7 +69,14 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import ( + Command, + Interrupt, + PregelTask, + Send, + StreamWriter, + interrupt, +) from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, @@ -262,8 +269,10 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: nonlocal tool_two_node_count tool_two_node_count += 1 if s["market"] == "DE": - raise NodeInterrupt("Just because...") - return {"my_key": " all good"} + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} tool_two_graph = StateGraph(State) tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) @@ -296,6 +305,25 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: with pytest.raises(ValueError, match="thread_id"): await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value ⛰️", "market": "DE"}, thread2 + ) + ] == [ + {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + ] + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear thread1 = {"configurable": {"thread_id": "1"}} # stop when about to enter node assert [ @@ -350,7 +378,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: # clear the interrupt and next tasks await tool_two.aupdate_state(thread1, None) - # interrupt is cleared, task will still run next + # interrupt is cleared, as well as the next tasks tup = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -376,7 +404,7 @@ async def test_node_not_cancelled_on_other_node_interrupted( checkpointer_name: str, ) -> None: class State(TypedDict): - hello: str + hello: Annotated[str, operator.add] awhiles = 0 inner_task_cancelled = False @@ -387,15 +415,14 @@ async def test_node_not_cancelled_on_other_node_interrupted( awhiles += 1 try: await asyncio.sleep(1) - return {"hello": "again"} + 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") + return {"hello": interrupt("I am bad")} builder = StateGraph(State) builder.add_node("agent", awhile) @@ -407,20 +434,25 @@ async def test_node_not_cancelled_on_other_node_interrupted( thread = {"configurable": {"thread_id": "1"}} # writes from "awhile" are applied to last chunk - assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "again"} + assert await graph.ainvoke({"hello": "world"}, thread) == { + "hello": "world again" + } assert not inner_task_cancelled assert awhiles == 1 - assert await graph.ainvoke(None, thread, debug=True) == {"hello": "again"} + assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world again"} assert not inner_task_cancelled assert awhiles == 1 - assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} + # resume with answer + assert await graph.ainvoke(Command(resume=" okay"), thread) == { + "hello": "world again okay" + } assert not inner_task_cancelled - assert awhiles == 2 + assert awhiles == 1 @pytest.mark.repeat(10) From 62d3a85b0707596afa6a6b36719324973c489f1b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:46:54 -0800 Subject: [PATCH 09/22] Add sync test --- libs/langgraph/tests/test_pregel.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bce736ac5..d801d9c98 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -79,7 +79,14 @@ from langgraph.pregel import ( from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import ( + Command, + Interrupt, + PregelTask, + Send, + StreamWriter, + interrupt, +) from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, @@ -8360,8 +8367,10 @@ def test_dynamic_interrupt( nonlocal tool_two_node_count tool_two_node_count += 1 if s["market"] == "DE": - raise NodeInterrupt("Just because...") - return {"my_key": " all good"} + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} tool_two_graph = StateGraph(State) tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) @@ -8393,6 +8402,20 @@ def test_dynamic_interrupt( with pytest.raises(ValueError, match="thread_id"): tool_two.invoke({"my_key": "value", "market": "DE"}) + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks thread1 = {"configurable": {"thread_id": "1"}} # stop when about to enter node assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { From c83b8f6d04a9791b65ab47f1d5305ad062fc8913 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:49:58 -0800 Subject: [PATCH 10/22] Update --- libs/scheduler-kafka/tests/test_subgraph.py | 18 ++++++++++++------ .../tests/test_subgraph_sync.py | 18 ++++++++++++------ 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index fd2530843..ebaaea580 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -194,8 +194,9 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -258,8 +259,9 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -352,8 +354,9 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -456,8 +459,9 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -515,8 +519,9 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -630,8 +635,9 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index 32d0ceea0..75b9d6e73 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -193,8 +193,9 @@ def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -255,10 +256,11 @@ def test_subgraph_w_interrupt( "__pregel_read": None, "__pregel_send": None, "__pregel_ensure_latest": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -350,9 +352,10 @@ def test_subgraph_w_interrupt( "__pregel_send": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -453,9 +456,10 @@ def test_subgraph_w_interrupt( "__pregel_send": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_resuming": True, "__pregel_task_id": history[1].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -512,9 +516,10 @@ def test_subgraph_w_interrupt( "__pregel_send": None, "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_resuming": True, "__pregel_task_id": history[1].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -628,8 +633,9 @@ def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, - '__pregel_store': None, + "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, + "__pregel_resume_value": None, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] From 311e16dffd7b71deb8ed2fda66633fcdb164ef6d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:52:18 -0800 Subject: [PATCH 11/22] Remove prints --- libs/langgraph/langgraph/pregel/algo.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 79ddb8722..564c53022 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -500,7 +500,6 @@ def prepare_single_task( return packet = writes_for_path[task_path_t[2]][2] if not isinstance(packet, Send): - print("packet", task_path_t, writes_for_path) logger.warning( f"Ignoring invalid packet type {type(packet)} in pending writes" ) @@ -613,19 +612,6 @@ def prepare_single_task( if name not in processes: return proc = processes[name] - print( - "preparing task", - task_path, - pending_writes, - sorted( - (chan, read_channel(channels, chan, return_exception=True)) - for chan in proc.triggers - # if not isinstance( - # read_channel(channels, chan, return_exception=True), - # EmptyChannelError, - # ) - ), - ) version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) null_version = version_type() # type: ignore[misc] if null_version is None: @@ -666,7 +652,6 @@ def prepare_single_task( "langgraph_path": task_path, "langgraph_checkpoint_ns": task_checkpoint_ns, } - print("preparing task", task_id, task_path, pending_writes) if task_id_checksum is not None: assert task_id == task_id_checksum if for_execution: From ef3a1ee9979b5ba87408ae7a7bc9af0d41d4db67 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:54:04 -0800 Subject: [PATCH 12/22] Undo --- libs/langgraph/langgraph/pregel/loop.py | 31 +++++++++++-------------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 36811d02a..6a9b6a95e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -417,7 +417,6 @@ class PregelLoop(LoopProtocol): ) for key, values in mv_writes.items(): self._update_mv(key, values) - print("applied null writes", null_writes) # prepare next tasks self.tasks = prepare_next_tasks( self.checkpoint, @@ -544,7 +543,6 @@ class PregelLoop(LoopProtocol): # save writes for tid, ws in writes.items(): self.put_writes(tid, ws) - print("applied cmd", writes) # map inputs to channel updates elif input_writes := deque(map_input(input_keys, self.input)): # TODO shouldn't these writes be passed to put_writes too? @@ -559,22 +557,19 @@ class PregelLoop(LoopProtocol): } ) # discard any unfinished tasks from previous checkpoint - if not isinstance(self.input, Command): - discard_tasks = prepare_next_tasks( - self.checkpoint, - self.checkpoint_pending_writes, - self.nodes, - self.channels, - self.managed, - self.config, - self.step, - for_execution=True, - store=None, - checkpointer=None, - manager=None, - ) - else: - discard_tasks = {} + discard_tasks = prepare_next_tasks( + self.checkpoint, + self.checkpoint_pending_writes, + self.nodes, + self.channels, + self.managed, + self.config, + self.step, + for_execution=True, + store=None, + checkpointer=None, + manager=None, + ) # apply input writes mv_writes = apply_writes( self.checkpoint, From 87fc519ce749c83c772c822a57fad9e248f22467 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:54:56 -0800 Subject: [PATCH 13/22] Remove print --- libs/langgraph/langgraph/types.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3337fc601..d867d60be 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -317,7 +317,6 @@ def interrupt(value: Any) -> Any: from langgraph.utils.config import get_configurable conf = get_configurable() - print("interrupt", conf.get(CONFIG_KEY_RESUME_VALUE)) if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING: return resume else: From b3a4eaa9678d2371381ca2cf871b536d0219e0b8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:56:10 -0800 Subject: [PATCH 14/22] Remove print --- libs/langgraph/langgraph/pregel/io.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 58a593645..7c26731f1 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -22,8 +22,7 @@ from langgraph.types import Command, PregelExecutableTask, Send def is_task_id(task_id: str) -> bool: """Check if a string is a valid task id.""" try: - u = UUID(task_id) - print(u.version) + UUID(task_id) except ValueError: return False return True From 86d2847dabf968aa9c759044d31fee00593d4210 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 17:57:45 -0800 Subject: [PATCH 15/22] Use neg idx --- libs/checkpoint/langgraph/checkpoint/base/__init__.py | 3 ++- libs/checkpoint/langgraph/checkpoint/serde/types.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index a63bbce28..a3ad83b1a 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -25,6 +25,7 @@ from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( ERROR, INTERRUPT, + RESUME, SCHEDULED, ChannelProtocol, SendProtocol, @@ -450,4 +451,4 @@ Special writes (e.g. errors) map to negative indices, to avoid those writes from conflicting with regular writes. Each Checkpointer implementation should use this mapping in put_writes. """ -WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2, INTERRUPT: -3} +WRITES_IDX_MAP = {ERROR: -1, SCHEDULED: -2, INTERRUPT: -3, RESUME: -4} diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index 9286e9b19..e258735bc 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -13,6 +13,7 @@ from typing_extensions import Self ERROR = "__error__" SCHEDULED = "__scheduled__" INTERRUPT = "__interrupt__" +RESUME = "__resume__" TASKS = "__pregel_tasks" Value = TypeVar("Value", covariant=True) From 0d5c6201d3b20a7135018590798a97ebea83aa6a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 18:09:58 -0800 Subject: [PATCH 16/22] Disable in py 3.10 or below for async --- libs/langgraph/langgraph/utils/config.py | 10 ++++++++++ .../tests/__snapshots__/test_pregel_async.ambr | 10 +++------- libs/langgraph/tests/test_pregel_async.py | 8 ++++++++ 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index 064a2b408..47de1f021 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -1,3 +1,5 @@ +import asyncio +import sys from collections import ChainMap from typing import Any, Optional, Sequence @@ -293,6 +295,14 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: def get_configurable() -> dict[str, Any]: + if sys.version_info < (3, 11): + try: + if asyncio.current_task(): + raise RuntimeError( + "Python 3.11 or later required to use this in an async context" + ) + except RuntimeError: + pass if var_config := var_child_runnable_config.get(): return var_config[CONF] else: diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index 3d4021111..adef94c62 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -1334,18 +1334,14 @@ __start__([

__start__

]):::first router_node(router_node) normal_llm_node(normal_llm_node) - weather_graph_model_node(model_node) - weather_graph_weather_node(weather_node
__interrupt = before) + weather_graph(weather_graph) __end__([

__end__

]):::last __start__ --> router_node; normal_llm_node --> __end__; - weather_graph_weather_node --> __end__; + weather_graph --> __end__; router_node -.-> normal_llm_node; - router_node -.-> weather_graph_model_node; + router_node -.-> weather_graph; router_node -.-> __end__; - subgraph weather_graph - weather_graph_model_node --> weather_graph_weather_node; - end classDef default fill:#f2f0ff,line-height:1.2 classDef first fill-opacity:0 classDef last fill:#bfb6fc diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1e4823d7e..022076fd9 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -257,6 +257,10 @@ async def test_node_cancellation_on_other_node_exception_two() -> None: await graph.ainvoke(1) +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict): @@ -399,6 +403,10 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: ) +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( checkpointer_name: str, From 00964b18f6502d3588bb4db5cf6aa1e12b69891c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 18:12:25 -0800 Subject: [PATCH 17/22] Undo --- .../tests/__snapshots__/test_pregel_async.ambr | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index adef94c62..3d4021111 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -1334,14 +1334,18 @@ __start__([

__start__

]):::first router_node(router_node) normal_llm_node(normal_llm_node) - weather_graph(weather_graph) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) __end__([

__end__

]):::last __start__ --> router_node; normal_llm_node --> __end__; - weather_graph --> __end__; + weather_graph_weather_node --> __end__; router_node -.-> normal_llm_node; - router_node -.-> weather_graph; + router_node -.-> weather_graph_model_node; router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end classDef default fill:#f2f0ff,line-height:1.2 classDef first fill-opacity:0 classDef last fill:#bfb6fc From 16bfa80b587f2f3ecb41c227713e839c32bc2e77 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 12 Nov 2024 16:18:23 -0800 Subject: [PATCH 18/22] Make Command accept generic arg for destinations --- libs/langgraph/langgraph/graph/state.py | 11 +-- libs/langgraph/langgraph/types.py | 8 +- .../tests/__snapshots__/test_pregel.ambr | 75 +++++++++++++++++++ .../__snapshots__/test_pregel_async.ambr | 75 +++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 3 +- libs/langgraph/tests/test_pregel_async.py | 7 +- 6 files changed, 168 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index d3b95d9e2..c5b0cd958 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -14,7 +14,6 @@ from typing import ( Optional, Sequence, Type, - TypeVar, Union, cast, get_args, @@ -50,15 +49,13 @@ from langgraph.managed.base import ( from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, Command, RetryPolicy +from langgraph.types import All, Checkpointer, Command, N, RetryPolicy from langgraph.utils.fields import get_field_default from langgraph.utils.pydantic import create_model from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable logger = logging.getLogger(__name__) -N = TypeVar("N") - def _warn_invalid_state_schema(schema: Union[Type[Any], Any]) -> None: if isinstance(schema, type): @@ -81,7 +78,7 @@ def _get_node_name(node: RunnableLike) -> str: raise TypeError(f"Unsupported node type: {type(node)}") -class GraphCommand(Command, Generic[N]): +class GraphCommand(Generic[N], Command[N]): """One or more commands to update a StateGraph's state and go to, or send messages to nodes.""" __slots__ = ("goto",) @@ -90,9 +87,9 @@ class GraphCommand(Command, Generic[N]): self, *, update: Optional[dict[str, Any]] = None, - goto: Union[str, Sequence[str]] = (), send: Union[Send, Sequence[Send]] = (), resume: Optional[Union[Any, dict[str, Any]]] = None, + goto: Union[str, Sequence[str]] = (), ) -> None: super().__init__(update=update, send=send, resume=resume) self.goto = goto @@ -390,7 +387,7 @@ class StateGraph(Graph): input = input_hint if ( (rtn := hints.get("return")) - and get_origin(rtn) is GraphCommand + and get_origin(rtn) in (Command, GraphCommand) and (rargs := get_args(rtn)) and get_origin(rargs[0]) is Literal and (vals := get_args(rargs[0])) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index d867d60be..b910616a2 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -4,11 +4,14 @@ from typing import ( TYPE_CHECKING, Any, Callable, + Generic, + Hashable, Literal, NamedTuple, Optional, Sequence, Type, + TypeVar, Union, cast, ) @@ -221,7 +224,10 @@ class Send: ) -class Command: +N = TypeVar("N", bound=Hashable) + + +class Command(Generic[N]): """One or more commands to update the graph's state and send messages to nodes.""" __slots__ = ("update", "send", "resume") diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 4268a310a..1cfdd9b76 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -5108,6 +5108,81 @@ ''' # --- +# name: test_send_react_interrupt_control[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_simple_multi_edge ''' graph TD; diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index 3d4021111..46916c7a4 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -1302,6 +1302,81 @@ +---------+ ''' # --- +# name: test_send_react_interrupt_control[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([

__start__

]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_weather_subgraph[duckdb_aio] ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d801d9c98..81d45f4ee 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2691,7 +2691,7 @@ def test_send_react_interrupt( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_react_interrupt_control( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion ) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage @@ -2721,6 +2721,7 @@ def test_send_react_interrupt_control( builder.add_node(foo) builder.add_edge(START, "agent") graph = builder.compile() + assert graph.get_graph().draw_mermaid() == snapshot assert graph.invoke({"messages": [HumanMessage("hello")]}) == { "messages": [ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 022076fd9..df85ed650 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2973,7 +2973,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_react_interrupt_control(checkpointer_name: str) -> None: +async def test_send_react_interrupt_control( + checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage ai_message = AIMessage( @@ -2982,7 +2984,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], ) - async def agent(state) -> GraphCommand[Literal["foo"]]: + async def agent(state) -> Command[Literal["foo"]]: return GraphCommand( update={"messages": ai_message}, send=[Send(call["name"], call) for call in ai_message.tool_calls], @@ -3000,6 +3002,7 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: builder.add_node(foo) builder.add_edge(START, "agent") graph = builder.compile() + assert graph.get_graph().draw_mermaid() == snapshot assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { "messages": [ From 7fe6f888760e82b32623b8f5c95e7ef37131018e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Nov 2024 12:51:55 -0800 Subject: [PATCH 19/22] Add resumeable/ns properties to Interrupt --- libs/langgraph/langgraph/errors.py | 2 +- libs/langgraph/langgraph/types.py | 29 +++++++++++++++++++---- libs/langgraph/tests/test_pregel.py | 18 ++++++++++++-- libs/langgraph/tests/test_pregel_async.py | 28 +++++++++++++++++++--- 4 files changed, 67 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 2e3d13120..2450b42b1 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -70,7 +70,7 @@ class NodeInterrupt(GraphInterrupt): """Raised by a node to interrupt execution.""" def __init__(self, value: Any) -> None: - super().__init__([Interrupt(value)]) + super().__init__([Interrupt(value=value)]) class GraphDelegate(Exception): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index b910616a2..167de66ab 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,5 +1,6 @@ from collections import deque from dataclasses import dataclass +import sys from typing import ( TYPE_CHECKING, Any, @@ -47,6 +48,11 @@ StreamWriter = Callable[[Any], None] Always injected into nodes if requested as a keyword argument, but it's a no-op when not using stream_mode="custom".""" +if sys.version_info >= (3, 10): + _DC_KWARGS = {"kw_only": True, "slots": True} +else: + _DC_KWARGS = {} + def default_retry_on(exc: Exception) -> bool: import httpx @@ -104,9 +110,11 @@ class CachePolicy(NamedTuple): pass -@dataclass +@dataclass(**_DC_KWARGS) class Interrupt: value: Any + resumable: bool = False + ns: Optional[str] = None when: Literal["during"] = "during" @@ -318,12 +326,25 @@ class LoopProtocol: def interrupt(value: Any) -> Any: - from langgraph.constants import CONFIG_KEY_RESUME_VALUE, MISSING - from langgraph.errors import NodeInterrupt + from langgraph.constants import ( + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_RESUME_VALUE, + MISSING, + NS_SEP, + ) + from langgraph.errors import GraphInterrupt from langgraph.utils.config import get_configurable conf = get_configurable() if (resume := conf.get(CONFIG_KEY_RESUME_VALUE, MISSING)) and resume is not MISSING: return resume else: - raise NodeInterrupt(value) + raise GraphInterrupt( + ( + Interrupt( + value=value, + resumable=True, + ns=cast(str, conf[CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP), + ), + ) + ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 81d45f4ee..2f88b2437 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8409,7 +8409,15 @@ def test_dynamic_interrupt( assert [ c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] # resume with answer assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ @@ -8447,7 +8455,13 @@ def test_dynamic_interrupt( AnyStr(), "tool_two", (PULL, "tool_two"), - interrupts=(Interrupt("Just because..."),), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), ), ), config=tool_two.checkpointer.get_tuple(thread1).config, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index df85ed650..48599b63f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -318,7 +318,15 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: {"my_key": "value ⛰️", "market": "DE"}, thread2 ) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] # resume with answer assert [ @@ -336,7 +344,15 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: {"my_key": "value ⛰️", "market": "DE"}, thread1 ) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ { @@ -363,7 +379,13 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: AnyStr(), "tool_two", (PULL, "tool_two"), - interrupts=(Interrupt("Just because..."),), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), ), ), config=tup.config, From 03bc9ba6e631727675bc4f12c5cc01a3939c0961 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Nov 2024 13:11:28 -0800 Subject: [PATCH 20/22] Make Command a dataclass --- .../langgraph/checkpoint/serde/jsonplus.py | 18 +------ .../langgraph/checkpoint/serde/types.py | 9 ---- libs/langgraph/langgraph/graph/state.py | 24 ++++----- libs/langgraph/langgraph/types.py | 50 ++++--------------- libs/langgraph/tests/test_pregel.py | 8 +-- libs/langgraph/tests/test_pregel_async.py | 5 +- 6 files changed, 30 insertions(+), 84 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 670e85b3d..f8d280b96 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -25,7 +25,7 @@ from langchain_core.load.serializable import Serializable from zoneinfo import ZoneInfo from langgraph.checkpoint.serde.base import SerializerProtocol -from langgraph.checkpoint.serde.types import CommandProtocol, SendProtocol +from langgraph.checkpoint.serde.types import SendProtocol from langgraph.store.base import Item LC_REVIVER = Reviver() @@ -122,11 +122,6 @@ class JsonPlusSerializer(SerializerProtocol): return self._encode_constructor_args( obj.__class__, kwargs={"node": obj.node, "arg": obj.arg} ) - elif isinstance(obj, CommandProtocol): - return self._encode_constructor_args( - obj.__class__, - kwargs={k: getattr(obj, k) for k in obj.__all_slots__}, - ) elif isinstance(obj, (bytes, bytearray)): return self._encode_constructor_args( obj.__class__, method="fromhex", args=(obj.hex(),) @@ -407,17 +402,6 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]: (obj.__class__.__module__, obj.__class__.__name__, (obj.node, obj.arg)), ), ) - elif isinstance(obj, CommandProtocol): - return msgpack.ExtType( - EXT_CONSTRUCTOR_KW_ARGS, - _msgpack_enc( - ( - obj.__class__.__module__, - obj.__class__.__name__, - {k: getattr(obj, k) for k in obj.__all_slots__}, - ), - ), - ) elif dataclasses.is_dataclass(obj): # doesn't use dataclasses.asdict to avoid deepcopy and recursion return msgpack.ExtType( diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index e258735bc..1df967b5f 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -4,7 +4,6 @@ from typing import ( Protocol, Sequence, TypeVar, - Union, runtime_checkable, ) @@ -51,11 +50,3 @@ class SendProtocol(Protocol): def __repr__(self) -> str: ... def __eq__(self, value: object) -> bool: ... - - -@runtime_checkable -class CommandProtocol(Protocol): - # Mirrors langgraph.types.Command - update: Optional[dict[str, Any]] - send: Union[Any, Sequence[Any]] - __all_slots__: set[str] diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index c5b0cd958..c581d2259 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1,3 +1,4 @@ +import dataclasses import inspect import logging import typing @@ -49,7 +50,7 @@ from langgraph.managed.base import ( from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import All, Checkpointer, Command, N, RetryPolicy +from langgraph.types import _DC_KWARGS, All, Checkpointer, Command, N, RetryPolicy from langgraph.utils.fields import get_field_default from langgraph.utils.pydantic import create_model from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable @@ -78,21 +79,20 @@ def _get_node_name(node: RunnableLike) -> str: raise TypeError(f"Unsupported node type: {type(node)}") +@dataclasses.dataclass(**_DC_KWARGS) class GraphCommand(Generic[N], Command[N]): """One or more commands to update a StateGraph's state and go to, or send messages to nodes.""" - __slots__ = ("goto",) + goto: Union[str, Sequence[str]] = () - def __init__( - self, - *, - update: Optional[dict[str, Any]] = None, - send: Union[Send, Sequence[Send]] = (), - resume: Optional[Union[Any, dict[str, Any]]] = None, - goto: Union[str, Sequence[str]] = (), - ) -> None: - super().__init__(update=update, send=send, resume=resume) - self.goto = goto + def __repr__(self) -> str: + # get all non-None values + contents = ", ".join( + f"{key}={value!r}" + for key, value in dataclasses.asdict(self).items() + if value + ) + return f"Command({contents})" class StateNodeSpec(NamedTuple): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 167de66ab..86c192c0c 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,5 +1,5 @@ from collections import deque -from dataclasses import dataclass +import dataclasses import sys from typing import ( TYPE_CHECKING, @@ -49,9 +49,9 @@ Always injected into nodes if requested as a keyword argument, but it's a no-op when not using stream_mode="custom".""" if sys.version_info >= (3, 10): - _DC_KWARGS = {"kw_only": True, "slots": True} + _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True} else: - _DC_KWARGS = {} + _DC_KWARGS = {"frozen": True} def default_retry_on(exc: Exception) -> bool: @@ -110,7 +110,7 @@ class CachePolicy(NamedTuple): pass -@dataclass(**_DC_KWARGS) +@dataclasses.dataclass(**_DC_KWARGS) class Interrupt: value: Any resumable: bool = False @@ -235,53 +235,23 @@ class Send: N = TypeVar("N", bound=Hashable) +@dataclasses.dataclass(**_DC_KWARGS) class Command(Generic[N]): """One or more commands to update the graph's state and send messages to nodes.""" - __slots__ = ("update", "send", "resume") - - def __init__( - self, - *, - update: Optional[dict[str, Any]] = None, - send: Union[Send, Sequence[Send]] = (), - resume: Optional[Union[Any, dict[str, Any]]] = None, - ) -> None: - self.update = update - self.send = send - self.resume = resume - - @property - def __all_slots__(self) -> set[str]: - # get all slots from mro - slots = set() - for cls in type(self).__mro__: - if ss := getattr(cls, "__slots__", ()): - if isinstance(ss, str): - slots.add(ss) - else: - slots.update(ss) - return slots + update: Optional[dict[str, Any]] = None + send: Union[Send, Sequence[Send]] = () + resume: Optional[Union[Any, dict[str, Any]]] = None def __repr__(self) -> str: # get all non-None values contents = ", ".join( f"{key}={value!r}" - for key in self.__all_slots__ - if (value := getattr(self, key)) + for key, value in dataclasses.asdict(self).items() + if value ) return f"Command({contents})" - def __eq__(self, value: Any) -> bool: - return type(value) is type(self) and all( - getattr(self, key) == getattr(value, key) for key in self.__all_slots__ - ) - - def copy(self, **kwargs: Any) -> Self: - for slot in self.__all_slots__: - kwargs.setdefault(slot, getattr(self, slot)) - return self.__class__(**kwargs) - StreamChunk = tuple[tuple[str, ...], str, Any] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2f88b2437..e29476c19 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8,6 +8,7 @@ import warnings from collections import Counter from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from dataclasses import replace from random import randrange from typing import ( Annotated, @@ -1834,9 +1835,8 @@ def test_send_sequences() -> None: if isinstance(state, list) else ["|".join((self.name, str(state)))] ) - if isinstance(state, GraphCommand): - state.update = update - return state + if isinstance(state, Command): + return replace(state, update=update) else: return update @@ -1918,7 +1918,7 @@ def test_send_dedupe_on_resume( else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): - return state.copy(update=update) + return replace(state, update=update) else: return update diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 48599b63f..1469c018e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6,6 +6,7 @@ import sys import uuid from collections import Counter from contextlib import asynccontextmanager, contextmanager +from dataclasses import replace from time import perf_counter from typing import ( Annotated, @@ -2133,7 +2134,7 @@ async def test_send_sequences(checkpointer_name: str) -> None: else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): - return state.copy(update=update) + return replace(state, update=update) else: return update @@ -2237,7 +2238,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: else ["|".join((self.name, str(state)))] ) if isinstance(state, GraphCommand): - return state.copy(update=update) + return replace(state, update=update) else: return update From 9fd152ef3aec16c6fe1a40a2cb7efee112dc7678 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Nov 2024 13:14:03 -0800 Subject: [PATCH 21/22] format --- libs/langgraph/langgraph/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 86c192c0c..26389ed37 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,6 +1,6 @@ -from collections import deque import dataclasses import sys +from collections import deque from typing import ( TYPE_CHECKING, Any, From a94902db8a86560baa298fe99655e051a3edb0d4 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Nov 2024 13:17:02 -0800 Subject: [PATCH 22/22] Lint --- libs/langgraph/langgraph/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 26389ed37..104412d8e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -114,7 +114,7 @@ class CachePolicy(NamedTuple): class Interrupt: value: Any resumable: bool = False - ns: Optional[str] = None + ns: Optional[Sequence[str]] = None when: Literal["during"] = "during"