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)