diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py index 2492551e2..c5fa373f3 100644 --- a/libs/langgraph/langgraph/channels/manager.py +++ b/libs/langgraph/langgraph/channels/manager.py @@ -1,6 +1,6 @@ from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager from datetime import datetime, timezone -from typing import Any, AsyncGenerator, Generator, Mapping +from typing import Any, AsyncGenerator, Generator, Mapping, Optional from langchain_core.runnables import RunnableConfig @@ -43,7 +43,11 @@ async def AsyncChannelsManager( def create_checkpoint( - checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], step: int + checkpoint: Checkpoint, + channels: Mapping[str, BaseChannel], + step: int, + *, + id: Optional[str] = None, ) -> Checkpoint: """Create a checkpoint for the given channels.""" ts = datetime.now(timezone.utc).isoformat() @@ -56,7 +60,7 @@ def create_checkpoint( return Checkpoint( v=1, ts=ts, - id=str(uuid6(clock_seq=step)), + id=id or str(uuid6(clock_seq=step)), channel_values=values, channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 6d85fd188..72b8c93db 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -44,7 +44,7 @@ class MemorySaver(BaseCheckpointSaver): asyncio.run(coro) # Output: 2 """ - storage: defaultdict[str, dict[str, tuple[bytes, bytes]]] + storage: defaultdict[str, dict[str, tuple[bytes, bytes, Optional[str]]]] def __init__( self, @@ -72,7 +72,7 @@ class MemorySaver(BaseCheckpointSaver): thread_id = config["configurable"]["thread_id"] if ts := config["configurable"].get("thread_ts"): if saved := self.storage[thread_id].get(ts): - checkpoint, metadata = saved + checkpoint, metadata, parent_ts = saved writes = self.writes[(thread_id, ts)] return CheckpointTuple( config=config, @@ -81,11 +81,19 @@ class MemorySaver(BaseCheckpointSaver): pending_writes=[ (id, c, self.serde.loads(v)) for id, c, v in writes ], + parent_config={ + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None, ) else: if checkpoints := self.storage[thread_id]: ts = max(checkpoints.keys()) - checkpoint, metadata = checkpoints[ts] + checkpoint, metadata, parent_ts = checkpoints[ts] writes = self.writes[(thread_id, ts)] return CheckpointTuple( config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, @@ -94,6 +102,14 @@ class MemorySaver(BaseCheckpointSaver): pending_writes=[ (id, c, self.serde.loads(v)) for id, c, v in writes ], + parent_config={ + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None, ) def list( @@ -120,7 +136,9 @@ class MemorySaver(BaseCheckpointSaver): """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage for thread_id in thread_ids: - for ts, (checkpoint, metadata_b) in self.storage[thread_id].items(): + for ts, (checkpoint, metadata_b, parent_ts) in sorted( + self.storage[thread_id].items(), key=lambda x: x[0], reverse=True + ): # filter by thread_ts if before and ts >= before["configurable"]["thread_ts"]: continue @@ -143,6 +161,14 @@ class MemorySaver(BaseCheckpointSaver): config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, checkpoint=self.serde.loads(checkpoint), metadata=metadata, + parent_config={ + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None, ) def put( @@ -169,6 +195,7 @@ class MemorySaver(BaseCheckpointSaver): checkpoint["id"]: ( self.serde.dumps(checkpoint), self.serde.dumps(metadata), + config["configurable"].get("thread_ts"), # parent ) } ) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 4a50e8833..f3aeb6a2e 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -3,9 +3,19 @@ from typing import Any INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" +CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" +CONFIG_KEY_RESUMING = "__pregel_resuming" INTERRUPT = "__interrupt__" TASKS = "__pregel_tasks" -RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT} +RESERVED = { + INTERRUPT, + TASKS, + CONFIG_KEY_SEND, + CONFIG_KEY_READ, + CONFIG_KEY_CHECKPOINTER, + CONFIG_KEY_RESUMING, + INPUT, +} TAG_HIDDEN = "langsmith:hidden" START = "__start__" diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 9191437d3..eed9cb7af 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -28,3 +28,15 @@ class InvalidUpdateError(Exception): """Raised when attempting to update a channel with an invalid sequence of updates.""" pass + + +class GraphInterrupt(Exception): + """Raised when a subgraph is interrupted.""" + + pass + + +class EmptyInputError(Exception): + """Raised when graph receives an empty input.""" + + pass diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 7defcdc70..c0edcacc3 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -63,7 +63,9 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.constants import ( + CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, + CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, ) @@ -281,7 +283,13 @@ class Pregel( ) ) # these are provided by the Pregel class - if spec.id not in [CONFIG_KEY_READ, CONFIG_KEY_SEND] + if spec.id + not in [ + CONFIG_KEY_READ, + CONFIG_KEY_SEND, + CONFIG_KEY_CHECKPOINTER, + CONFIG_KEY_RESUMING, + ] ] @property @@ -699,6 +707,7 @@ class Pregel( Union[str, Sequence[str]], Optional[Sequence[str]], Optional[Sequence[str]], + Optional[BaseCheckpointSaver], ]: debug = debug if debug is not None else self.debug if output_keys is None: @@ -710,15 +719,24 @@ class Pregel( stream_mode = stream_mode if stream_mode is not None else self.stream_mode if not isinstance(stream_mode, list): stream_mode = [stream_mode] - if config is not None and config.get("configurable", {}).get(CONFIG_KEY_READ): + if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None: # if being called as a node in another graph, always use values mode stream_mode = ["values"] + if config is not None and config.get("configurable", {}).get( + CONFIG_KEY_CHECKPOINTER + ): + checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][ + CONFIG_KEY_CHECKPOINTER + ] + else: + checkpointer = self.checkpointer return ( debug, stream_mode, output_keys, interrupt_before, interrupt_after, + checkpointer, ) def stream( @@ -820,6 +838,7 @@ class Pregel( output_keys, interrupt_before, interrupt_after, + checkpointer, ) = self._defaults( config, stream_mode=stream_mode, @@ -830,7 +849,7 @@ class Pregel( ) with SyncPregelLoop( - input, config=config, checkpointer=self.checkpointer, graph=self + input, config=config, checkpointer=checkpointer, graph=self ) as loop: # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates @@ -1063,6 +1082,7 @@ class Pregel( output_keys, interrupt_before, interrupt_after, + checkpointer, ) = self._defaults( config, stream_mode=stream_mode, @@ -1072,7 +1092,7 @@ class Pregel( debug=debug, ) async with AsyncPregelLoop( - input, config=config, checkpointer=self.checkpointer, graph=self + input, config=config, checkpointer=checkpointer, graph=self ) as loop: aioloop = asyncio.get_event_loop() # Similarly to Bulk Synchronous Parallel / Pregel model @@ -1201,6 +1221,7 @@ class Pregel( read_channels(loop.channels, output_keys) ) except BaseException as e: + # TODO use on_chain_end if exc is GraphInterrupt await asyncio.shield(run_manager.on_chain_error(e)) raise diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index e93d6e14a..13ce7e183 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -26,9 +26,11 @@ from langchain_core.runnables.config import ( from langgraph.channels.base import BaseChannel from langgraph.channels.context import Context from langgraph.channels.manager import ChannelsManager, create_checkpoint -from langgraph.checkpoint.base import Checkpoint, copy_checkpoint +from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, copy_checkpoint from langgraph.constants import ( + CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, + CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, RESERVED, @@ -213,6 +215,8 @@ def prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[False], + is_resuming: bool = False, + checkpointer: Literal[None] = None, manager: Literal[None] = None, ) -> list[PregelTaskDescription]: ... @@ -227,6 +231,8 @@ def prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[True], + is_resuming: bool, + checkpointer: Optional[BaseCheckpointSaver], manager: Union[None, ParentRunManager, AsyncParentRunManager], ) -> list[PregelExecutableTask]: ... @@ -241,6 +247,8 @@ def prepare_next_tasks( step: int, *, for_execution: bool, + is_resuming: bool = False, + checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, ) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]: tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] @@ -291,6 +299,8 @@ def prepare_next_tasks( PregelTaskWrites(packet.node, writes, triggers), config, ), + # in Send we can't checkpoint nested graphs + # as they could be running in parallel }, ), triggers, @@ -332,6 +342,12 @@ def prepare_next_tasks( "langgraph_task_idx": len(tasks), } task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) + if parent_thread_id := config.get("configurable", {}).get( + "thread_id" + ): + thread_id: Optional[str] = f"{parent_thread_id}-{name}" + else: + thread_id = None writes = deque() tasks.append( PregelExecutableTask( @@ -363,6 +379,10 @@ def prepare_next_tasks( PregelTaskWrites(name, writes, triggers), config, ), + CONFIG_KEY_CHECKPOINTER: checkpointer, + CONFIG_KEY_RESUMING: is_resuming, + "thread_id": thread_id, + "thread_ts": checkpoint["id"], }, ), triggers, diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 389c59e81..70c1e69b7 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -70,14 +70,15 @@ def map_debug_tasks( if config is not None and TAG_HIDDEN in config.get("tags", []): continue + metadata = config["metadata"].copy() + metadata.pop("thread_ts", None) + yield { "type": "task", "timestamp": ts, "step": step, "payload": { - "id": str( - uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"]))) - ), + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))), "name": name, "input": input, "triggers": triggers, @@ -95,14 +96,15 @@ def map_debug_task_results( if config is not None and TAG_HIDDEN in config.get("tags", []): continue + metadata = config["metadata"].copy() + metadata.pop("thread_ts", None) + yield { "type": "task_result", "timestamp": ts, "step": step, "payload": { - "id": str( - uuid5(TASK_NAMESPACE, json.dumps((name, step, config["metadata"]))) - ), + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))), "name": name, "result": [w for w in writes if w[0] in stream_channels_list], }, diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 352b8cf51..5da78a26c 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -1,14 +1,14 @@ import asyncio import concurrent.futures import sys -from contextlib import contextmanager +from contextlib import ExitStack from contextvars import copy_context from types import TracebackType from typing import ( AsyncContextManager, Awaitable, Callable, - Iterator, + ContextManager, Optional, Protocol, TypeVar, @@ -18,6 +18,8 @@ from langchain_core.runnables import RunnableConfig from langchain_core.runnables.config import get_executor_for_config from typing_extensions import ParamSpec +from langgraph.errors import GraphInterrupt + P = ParamSpec("P") T = TypeVar("T") @@ -34,41 +36,63 @@ class Submit(Protocol[P, T]): ... -@contextmanager -def BackgroundExecutor(config: RunnableConfig) -> Iterator[Submit]: - tasks: dict[concurrent.futures.Future, bool] = {} - with get_executor_for_config(config) as executor: +class BackgroundExecutor(ContextManager): + def __init__(self, config: RunnableConfig) -> None: + self.stack = ExitStack() + self.executor = self.stack.enter_context(get_executor_for_config(config)) + self.tasks: dict[concurrent.futures.Future, bool] = {} - def done(task: concurrent.futures.Future) -> None: - try: - task.result() - except BaseException: - pass - else: - tasks.pop(task) - - def submit( - fn: Callable[P, T], - *args: P.args, - __name__: Optional[str] = None, # currently not used in sync version - __cancel_on_exit__: bool = False, - **kwargs: P.kwargs, - ) -> concurrent.futures.Future: - task = executor.submit(fn, *args, **kwargs) - tasks[task] = __cancel_on_exit__ - task.add_done_callback(done) - return task + def submit( + self, + fn: Callable[P, T], + *args: P.args, + __name__: Optional[str] = None, # currently not used in sync version + __cancel_on_exit__: bool = False, + **kwargs: P.kwargs, + ) -> concurrent.futures.Future[T]: + task = self.executor.submit(fn, *args, **kwargs) + self.tasks[task] = __cancel_on_exit__ + task.add_done_callback(self.done) + return task + def done(self, task: concurrent.futures.Future) -> None: try: - yield submit - finally: - for task, cancel in tasks.items(): - if cancel: - task.cancel() - # executor waits for all tasks to finish on exit - for task in tasks: - # the first task to have raised an exception will be re-raised here - task.result() + task.result() + except GraphInterrupt: + # This exception is an interruption signal, not an error + # so we don't want to re-raise it on exit + self.tasks.pop(task) + except BaseException: + pass + else: + self.tasks.pop(task) + + def __enter__(self) -> "submit": + return self.submit + + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + # cancel all tasks that should be cancelled + for task, cancel in self.tasks.items(): + if cancel: + task.cancel() + # wait for all tasks to finish + if tasks := {t for t in self.tasks if not t.done()}: + concurrent.futures.wait(tasks) + # shutdown the executor + self.stack.__exit__(exc_type, exc_value, traceback) + # re-raise the first exception that occurred in a task + if exc_type is None: + # if there's already an exception being raised, don't raise another one + for task in self.tasks: + try: + task.result() + except concurrent.futures.CancelledError: + pass class AsyncBackgroundExecutor(AsyncContextManager): @@ -97,24 +121,39 @@ class AsyncBackgroundExecutor(AsyncContextManager): def done(self, task: asyncio.Task) -> None: try: task.result() + except GraphInterrupt: + # This exception is an interruption signal, not an error + # so we don't want to re-raise it on exit + self.tasks.pop(task) except BaseException: pass else: self.tasks.pop(task) - async def __aenter__(self) -> "submit": + async def __aenter__(self) -> Submit: return self.submit - async def exit(self) -> None: - fut = asyncio.gather(*self.tasks, return_exceptions=True) - try: - rtns = await asyncio.shield(fut) - finally: - del self.tasks - for rtn in rtns: - # if this is ever changed to BaseException, need to ignore CancelledError - if isinstance(rtn, Exception): - raise rtn + async def exit( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + # cancel all tasks that should be cancelled + for task, cancel in self.tasks.items(): + if cancel: + task.cancel(self.sentinel) + # wait for all tasks to finish + if self.tasks: + await asyncio.wait(self.tasks) + # re-raise the first exception that occurred in a task + if exc_type is None: + # if there's already an exception being raised, don't raise another one + for task in self.tasks: + try: + task.result() + except asyncio.CancelledError: + pass async def __aexit__( self, @@ -122,8 +161,8 @@ class AsyncBackgroundExecutor(AsyncContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - for task, cancel in self.tasks.items(): - if cancel: - task.cancel(self.sentinel) + # we cannot use `await` outside of asyncio.shield, as this code can run + # after owning task is cancelled, so pulling async logic to separate method + # wait for all background tasks to finish, shielded from cancellation - await asyncio.shield(self.exit()) + await asyncio.shield(self.exit(exc_type, exc_value, traceback)) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 73bb8f20f..f9c00c62a 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,4 +1,5 @@ import asyncio +import concurrent.futures from collections import deque from contextlib import AsyncExitStack, ExitStack from types import TracebackType @@ -38,7 +39,8 @@ from langgraph.checkpoint.base import ( copy_checkpoint, empty_checkpoint, ) -from langgraph.constants import INPUT, INTERRUPT +from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_RESUMING, INPUT, INTERRUPT +from langgraph.errors import EmptyInputError, GraphInterrupt from langgraph.managed.base import ( AsyncManagedValuesManager, ManagedValueMapping, @@ -66,6 +68,8 @@ if TYPE_CHECKING: V = TypeVar("V") INPUT_DONE = object() +INPUT_RESUMING = object() +EMPTY_SEQ = () class PregelLoop: @@ -76,8 +80,16 @@ class PregelLoop: checkpointer_put_writes: Optional[ Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any] ] - checkpointer_put: Optional[ - Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], Any] + _checkpointer_put_after_previous: Optional[ + Callable[ + [ + Optional[concurrent.futures.Future], + RunnableConfig, + Sequence[tuple[str, Any]], + str, + ], + Any, + ] ] graph: "Pregel" @@ -95,9 +107,27 @@ class PregelLoop: ] tasks: Sequence[PregelExecutableTask] stream: deque[Tuple[str, Any]] + is_nested: bool # public + def __init__( + self, + input: Optional[Any], + *, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", + ) -> None: + self.stream = deque() + self.input = input + self.config = config + self.checkpointer = checkpointer + self.graph = graph + # TODO if managed values no longer needs graph we can replace with + # managed_specs, channel_specs + self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) + def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: """Mark tasks as scheduled, to be used by queue-based executors.""" raise NotImplementedError @@ -122,9 +152,9 @@ class PregelLoop: def tick( self, *, - output_keys: Union[str, Sequence[str]] = None, - interrupt_after: Optional[Sequence[str]] = None, - interrupt_before: Optional[Sequence[str]] = None, + output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, + interrupt_after: Sequence[str] = EMPTY_SEQ, + interrupt_before: Sequence[str] = EMPTY_SEQ, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, ) -> bool: """Execute a single iteration of the Pregel loop. @@ -133,7 +163,7 @@ class PregelLoop: if self.status != "pending": raise RuntimeError("Cannot tick when status is no longer 'pending'") - if self.input is not INPUT_DONE: + if self.input not in (INPUT_DONE, INPUT_RESUMING): self._first() elif all(task.writes for task in self.tasks): writes = [w for t in self.tasks for w in t.writes] @@ -165,7 +195,10 @@ class PregelLoop: # after execution, check if we should interrupt if should_interrupt(self.checkpoint, interrupt_after, self.tasks): self.status = "interrupt_after" - return False + if self.is_nested: + raise GraphInterrupt(self) + else: + return False else: return False @@ -184,6 +217,8 @@ class PregelLoop: self.step, for_execution=True, manager=manager, + checkpointer=self.checkpointer, + is_resuming=self.input is INPUT_RESUMING, ) # if no more tasks, we're done @@ -199,12 +234,20 @@ class PregelLoop: # if all tasks have finished, re-tick if all(task.writes for task in self.tasks): - return self.tick() + return self.tick( + output_keys=output_keys, + interrupt_after=interrupt_after, + interrupt_before=interrupt_before, + manager=manager, + ) # before execution, check if we should interrupt if should_interrupt(self.checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" - return False + if self.is_nested: + raise GraphInterrupt() + else: + return False # produce debug output self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks)) @@ -214,8 +257,23 @@ class PregelLoop: # private def _first(self) -> None: + # resuming from previous checkpoint requires + # - finding a previous checkpoint + # - receiving None input (outer graph) or RESUMING flag (subgraph) + is_resuming = bool(self.checkpoint["channel_versions"]) and bool( + self.config.get("configurable", {}).get(CONFIG_KEY_RESUMING) + or self.input is None + ) + + # proceed past previous checkpoint + if is_resuming: + self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) + for k in self.channels: + if k in self.checkpoint["channel_versions"]: + version = self.checkpoint["channel_versions"][k] + self.checkpoint["versions_seen"][INTERRUPT][k] = version # map inputs to channel updates - if input_writes := deque(map_input(self.graph.input_channels, self.input)): + elif input_writes := deque(map_input(self.graph.input_channels, self.input)): # discard any unfinished tasks from previous checkpoint discard_tasks = prepare_next_tasks( self.checkpoint, @@ -237,31 +295,33 @@ class PregelLoop: # save input checkpoint self._put_checkpoint({"source": "input", "writes": self.input}) else: - # no input is taken as signal to proceed past previous interrupt - self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) - for k in self.channels: - if k in self.checkpoint["channel_versions"]: - version = self.checkpoint["channel_versions"][k] - self.checkpoint["versions_seen"][INTERRUPT][k] = version + raise EmptyInputError(f"Received no input for {self.graph.input_channels}") # done with input - self.input = INPUT_DONE + self.input = INPUT_RESUMING if is_resuming else INPUT_DONE - def _put_checkpoint( - self, - metadata: CheckpointMetadata, - ) -> None: + def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: # assign step metadata["step"] = self.step # bail if no checkpointer - if self.checkpointer_put is not None: + if self._checkpointer_put_after_previous is not None: # create new checkpoint self.checkpoint_metadata = metadata self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.step + self.checkpoint, + self.channels, + self.step, + # child graphs keep at most one checkpoint per parent checkpoint + # this is achieved by writing child checkpoints as progress is made + # (so that error recovery / resuming from interrupt don't lose work) + # but doing so always with an id equal to that of the parent checkpoint + id=self.config["configurable"]["thread_ts"] if self.is_nested else None, ) # save it, without blocking - self.submit( - self.checkpointer_put, + # if there's a previous checkpoint save in progress, wait for it + # ensuring checkpointers receive checkpoints in order + self._put_checkpoint_fut = self.submit( + self._checkpointer_put_after_previous, + getattr(self, "_put_checkpoint_fut", None), self.checkpoint_config, copy_checkpoint(self.checkpoint), self.checkpoint_metadata, @@ -287,6 +347,15 @@ class PregelLoop: # increment step self.step += 1 + def _suppress_interrupt( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + if exc_type is GraphInterrupt and not self.is_nested: + return True + class SyncPregelLoop(PregelLoop, ContextManager): def __init__( @@ -297,19 +366,29 @@ class SyncPregelLoop(PregelLoop, ContextManager): checkpointer: Optional[BaseCheckpointSaver], graph: "Pregel", ) -> None: - self.stream = deque() + super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) self.stack = ExitStack() - self.input = input - self.config = config - self.checkpointer = checkpointer - self.checkpointer_get_next_version = ( - checkpointer.get_next_version if checkpointer else increment - ) - self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None - self.checkpointer_put = checkpointer.put if checkpointer else None - self.graph = graph - # TODO if managed values no longer needs graph we can replace with - # managed_specs, channel_specs + self.stack.push(self._suppress_interrupt) + if checkpointer: + self.checkpointer_get_next_version = checkpointer.get_next_version + self.checkpointer_put_writes = checkpointer.put_writes + else: + self.checkpointer_get_next_version = increment + self._checkpointer_put_after_previous = None + self.checkpointer_put_writes = None + + def _checkpointer_put_after_previous( + self, + prev: Optional[concurrent.futures.Future], + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + ) -> RunnableConfig: + try: + if prev is not None: + prev.result() + finally: + self.checkpointer.put(config, checkpoint, metadata) # context manager @@ -349,6 +428,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # unwind stack del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) @@ -362,21 +442,29 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): checkpointer: Optional[BaseCheckpointSaver], graph: "Pregel", ) -> None: - self.stream = deque() + super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) self.stack = AsyncExitStack() - self.input = input - self.config = config - self.checkpointer = checkpointer - self.checkpointer_get_next_version = ( - checkpointer.get_next_version if checkpointer else increment - ) - self.checkpointer_put_writes = ( - checkpointer.aput_writes if checkpointer else None - ) - self.checkpointer_put = checkpointer.aput if checkpointer else None - self.graph = graph - # TODO if managed values no longer needs graph we can replace with - # managed_specs, channel_specs + self.stack.push(self._suppress_interrupt) + if checkpointer: + self.checkpointer_get_next_version = checkpointer.get_next_version + self.checkpointer_put_writes = checkpointer.aput_writes + else: + self.checkpointer_get_next_version = increment + self._checkpointer_put_after_previous = None + self.checkpointer_put_writes = None + + async def _checkpointer_put_after_previous( + self, + prev: Optional[asyncio.Task], + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + ) -> RunnableConfig: + try: + if prev is not None: + await prev + finally: + await self.checkpointer.aput(config, checkpoint, metadata) # context manager @@ -418,6 +506,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # unwind stack del self.graph return await asyncio.shield( self.stack.__aexit__(exc_type, exc_value, traceback) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 94075d612..0b7a1755f 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -2783,6 +2783,20 @@ pytest = ">=6.2.5" [package.extras] dev = ["pre-commit", "pytest-asyncio", "tox"] +[[package]] +name = "pytest-repeat" +version = "0.9.3" +description = "pytest plugin for repeating tests" +optional = false +python-versions = ">=3.7" +files = [ + {file = "pytest_repeat-0.9.3-py3-none-any.whl", hash = "sha256:26ab2df18226af9d5ce441c858f273121e92ff55f5bb311d25755b8d7abdd8ed"}, + {file = "pytest_repeat-0.9.3.tar.gz", hash = "sha256:ffd3836dfcd67bb270bec648b330e20be37d2966448c4148c4092d1e8aba8185"}, +] + +[package.dependencies] +pytest = "*" + [[package]] name = "pytest-watcher" version = "0.4.2" @@ -4165,4 +4179,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "0d877d3879473de43aca1e1d36a8f420ff3f4b140807cb5cc24935d9114947be" +content-hash = "18b26895b05f2f7cdcd08d59ba164ac788e0e8005e3ec5d7089469b4f3a96aea" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 593e659c8..b3d76de2a 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -32,6 +32,7 @@ langchain-openai = ">=0.1.2" langchain-anthropic = ">=0.1.8" dataclasses-json = "^0.6.7" pytest-xdist = {extras = ["psutil"], version = "^3.6.1"} +pytest-repeat = "^0.9.3" [tool.poetry.group.dev] optional = true @@ -62,7 +63,7 @@ omit = ["tests/*"] [tool.pytest-watcher] now = true delay = 0.1 -runner_args = ["-x", "--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"] +runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"] patterns = ["*.py"] [build-system] diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index db1212ae9..c0ede2f20 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -30,9 +30,11 @@ class MemorySaverAssertImmutable(MemorySaver): self, *, serde: Optional[SerializerProtocol] = None, + put_sleep: Optional[float] = None, ) -> None: super().__init__(serde=serde) self.storage_for_copies = defaultdict(dict) + self.put_sleep = put_sleep def put( self, @@ -40,6 +42,10 @@ class MemorySaverAssertImmutable(MemorySaver): checkpoint: Checkpoint, metadata: Optional[CheckpointMetadata] = None, ) -> None: + if self.put_sleep: + import time + + time.sleep(self.put_sleep) # assert checkpoint hasn't been modified since last written thread_id = config["configurable"]["thread_id"] if saved := super().get(config): @@ -85,7 +91,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): configurable = config["configurable"].copy() # remove thread_ts to make testing simpler - configurable.pop("thread_ts", None) + thread_ts = configurable.pop("thread_ts", None) self.storage[config["configurable"]["thread_id"]].update( { @@ -93,6 +99,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): self.serde.dumps(checkpoint), # merge configurable fields and metadata self.serde.dumps({**configurable, **metadata}), + thread_ts, ) } ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2646a64fd..db902eca0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8,6 +8,7 @@ from contextlib import contextmanager from typing import ( Annotated, Any, + Callable, Dict, Generator, List, @@ -606,35 +607,10 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: assert snapshot.next == () # list history - assert [c for c in app.get_state_history({"configurable": {"thread_id": 1}})] == [ + thread1 = {"configurable": {"thread_id": 1}} + assert [c for c in app.get_state_history(thread1)] == [ StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": -1, "writes": 2}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, + values={"inbox": 4, "output": 5, "input": 3}, next=(), config={ "configurable": { @@ -642,48 +618,9 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 6, "writes": 5}, created_at=AnyStr(), - metadata={"source": "loop", "step": 1, "writes": 4}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 2, "writes": 20}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 3, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 4, "writes": 3}, - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][1].config, ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, @@ -694,12 +631,51 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: "thread_ts": AnyStr(), } }, - created_at=AnyStr(), metadata={"source": "loop", "step": 5, "writes": None}, - parent_config=None, + created_at=AnyStr(), + parent_config=[*app.checkpointer.list(thread1)][2].config, ), StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=[*app.checkpointer.list(thread1)][3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=[*app.checkpointer.list(thread1)][4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=[*app.checkpointer.list(thread1)][5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, next=(), config={ "configurable": { @@ -707,8 +683,34 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=[*app.checkpointer.list(thread1)][6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=[*app.checkpointer.list(thread1)][7].config, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, created_at=AnyStr(), - metadata={"source": "loop", "step": 6, "writes": 5}, parent_config=None, ), ] @@ -1723,6 +1725,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert ( app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ @@ -1771,6 +1774,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1878,6 +1882,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # test state get/update methods with interrupt_before @@ -1928,6 +1933,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -1970,6 +1976,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2077,6 +2084,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # test re-invoke to continue with interrupt_before @@ -2127,6 +2135,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2526,6 +2535,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -2565,6 +2575,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: }, }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2636,6 +2647,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # test state get/update methods with interrupt_before @@ -2684,6 +2696,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -2723,6 +2736,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2794,6 +2808,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # test w interrupt before all @@ -2818,6 +2833,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2854,6 +2870,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2911,6 +2928,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2969,6 +2987,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -3026,6 +3045,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -4129,6 +4149,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -4182,6 +4203,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -4286,6 +4308,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: }, }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -4337,6 +4360,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -4620,6 +4644,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -4663,6 +4688,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -4746,6 +4772,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -4787,6 +4814,7 @@ def test_message_graph( "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt = workflow.compile( @@ -4850,6 +4878,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -4896,6 +4925,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -4979,6 +5009,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -5020,6 +5051,7 @@ def test_message_graph( "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # add an extra message as if it came from "tools" node @@ -5061,6 +5093,7 @@ def test_message_graph( "step": 6, "writes": {"tools": ("ai", "an extra message")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5342,6 +5375,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -5385,6 +5419,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -5468,6 +5503,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -5509,6 +5545,7 @@ def test_root_graph( "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt = workflow.compile( @@ -5572,6 +5609,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -5618,6 +5656,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -5701,6 +5740,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -5742,6 +5782,7 @@ def test_root_graph( "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # add an extra message as if it came from "tools" node @@ -5783,6 +5824,7 @@ def test_root_graph( "step": 6, "writes": {"tools": ("ai", "an extra message")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # create new graph with one more state key, reuse previous thread history @@ -5856,6 +5898,7 @@ def test_root_graph( "step": 6, "writes": {"tools": ("ai", "an extra message")}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # new input is merged to old state @@ -6885,6 +6928,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> "step": 4, "writes": {"retriever_one": {"docs": ["doc5"]}}, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ @@ -7524,6 +7568,1111 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.repeat(10) +@pytest.mark.parametrize( + "checkpointer_fct", + [ + lambda: MemorySaverAssertImmutable(put_sleep=0.2), + lambda: SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_nested_graph_interrupts( + checkpointer_fct: Callable[[], BaseCheckpointSaver], +) -> None: + try: + checkpointer = checkpointer_fct() + + class InnerState(TypedDict): + my_key: str + my_other_key: str + + def inner_1(state: InnerState): + return { + "my_key": state["my_key"] + " here", + "my_other_key": state["my_key"], + } + + def inner_2(state: InnerState): + return { + "my_key": state["my_key"] + " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + def outer_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def outer_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("outer_1", outer_1) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_2", outer_2) + graph.set_entry_point("outer_1") + graph.add_edge("outer_1", "inner") + graph.add_edge("inner", "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert app.invoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [*app.stream({"my_key": "my value"}, config)] == [ + {"outer_1": {"my_key": "hi my value"}}, + ] + assert [*app.stream(None, config)] == [ + {"inner": {"my_key": "hi my value here and there"}}, + {"outer_2": {"my_key": "hi my value here and there and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + + # test interrupts BEFORE the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) + config = {"configurable": {"thread_id": "4"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [*app.stream(None, config, stream_mode="values")] == [] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # test interrupts AFTER the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) + config = {"configurable": {"thread_id": "5"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + # interrupted after "inner" + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there and back again", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # test restarting from thread_ts + config = {"configurable": {"thread_id": "6"}} + app = graph.compile(checkpointer=checkpointer) + assert app.invoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value" + } + state_history = [c for c in app.get_state_history(config)] + assert state_history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + child_state_history = [ + c for c in app.get_state_history({"configurable": {"thread_id": "6-inner"}}) + ] + assert child_state_history == [ + StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + ), + # there should be a single child checkpoint because we only keep + # one child checkpoint per parent checkpoint (in which child ran) + ] + + # check that child snapshot matches id of parent + child_snapshot = child_state_history[0] + assert ( + child_snapshot.config["configurable"]["thread_ts"] + == state_history[0].config["configurable"]["thread_ts"] + ) + # check resuming from interrupt w/ thread_ts + interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] + before_interrupt_config = before_interrupt_state_snapshot.config + # going to get to interrupt again here, so the output is None + assert app.invoke(None, before_interrupt_config, debug=True) == { + "my_key": "hi my value" + } + # one more "identical" snapshot than before, at top of list + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # going to restart from interrupt + interrupt_config = interrupt_state_snapshot.config + assert app.invoke(None, interrupt_config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> None: + try: + + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + time.sleep(0.1) + return {"my_key": "got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def outer_1(state: State): + return {"my_key": " and parallel"} + + def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": ""}, config, debug=True) == { + "my_key": "", + } + + assert app.invoke(None, config, debug=True) == { + "my_key": "got here and there and parallel and back again", + } + + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [*app.stream({"my_key": ""}, config)] == [ + # we got to parallel node first + {"outer_1": {"my_key": " and parallel"}}, + ] + assert [*app.stream(None, config)] == [ + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + { + "my_key": "", + }, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "got here and there and parallel", + }, + { + "my_key": "got here and there and parallel and back again", + }, + ] + + # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + {"my_key": ""} + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [*app.stream(None, config, stream_mode="values")] == [] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "got here and there and parallel", + }, + { + "my_key": "got here and there and parallel and back again", + }, + ] + + # test interrupts AFTER the parallel node + 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": ""} + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + {"my_key": "got here and there and parallel"}, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "got here and there and parallel and back again", + }, + ] + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_doubly_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: + try: + + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node("child_1", grandchild.compile(interrupt_before=["grandchild_2"])) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } + + assert app.invoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [*app.stream({"my_key": "my value"}, config)] == [ + {"parent_1": {"my_key": "hi my value"}}, + ] + assert [*app.stream(None, config)] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, None) + + def test_repeat_condition(snapshot: SnapshotAssertion) -> None: class AgentState(TypedDict): hello: str @@ -7576,7 +8725,7 @@ def test_checkpoint_metadata() -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) - from langchain_core.messages import AIMessage, AnyMessage + from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage from langchain_core.prompts import ChatPromptTemplate from langchain_core.tools import tool @@ -7651,7 +8800,7 @@ def test_checkpoint_metadata() -> None: # assertions # invoke graph w/o interrupt - app.invoke( + assert app.invoke( {"messages": ["what is weather in sf"]}, { "configurable": { @@ -7660,7 +8809,30 @@ def test_checkpoint_metadata() -> None: "test_config_2": "bar", }, }, - ) + ) == { + "messages": [ + HumanMessage(content="what is weather in sf", id=AnyStr()), + AIMessage( + content="", + id=AnyStr(), + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ), + ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id=AnyStr()), + ] + } config = {"configurable": {"thread_id": "1"}} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0e927e518..c3f8bb3bf 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8,6 +8,7 @@ from typing import ( Any, AsyncGenerator, AsyncIterator, + Callable, Dict, Generator, List, @@ -723,37 +724,10 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N assert snapshot.next == () # list history - assert [ - c async for c in app.aget_state_history({"configurable": {"thread_id": 1}}) - ] == [ + thread1 = {"configurable": {"thread_id": 1}} + assert [c async for c in app.aget_state_history(thread1)] == [ StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": -1, "writes": 2}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, + values={"inbox": 4, "output": 5, "input": 3}, next=(), config={ "configurable": { @@ -761,48 +735,9 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 6, "writes": 5}, created_at=AnyStr(), - metadata={"source": "loop", "step": 1, "writes": 4}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 2, "writes": 20}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 3, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 4, "writes": 3}, - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][1].config, ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, @@ -813,12 +748,51 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N "thread_ts": AnyStr(), } }, - created_at=AnyStr(), metadata={"source": "loop", "step": 5, "writes": None}, - parent_config=None, + created_at=AnyStr(), + parent_config=[c async for c in app.checkpointer.alist(thread1)][2].config, ), StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=[c async for c in app.checkpointer.alist(thread1)][3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=[c async for c in app.checkpointer.alist(thread1)][4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=[c async for c in app.checkpointer.alist(thread1)][5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, next=(), config={ "configurable": { @@ -826,8 +800,34 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=[c async for c in app.checkpointer.alist(thread1)][6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=[c async for c in app.checkpointer.alist(thread1)][7].config, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, created_at=AnyStr(), - metadata={"source": "loop", "step": 6, "writes": 5}, parent_config=None, ), ] @@ -1936,6 +1936,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) await app_w_interrupt.aupdate_state( @@ -1980,6 +1983,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2089,6 +2095,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) # test state get/update methods with interrupt_before @@ -2144,6 +2153,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) await app_w_interrupt.aupdate_state( @@ -2188,6 +2200,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2297,6 +2312,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) # test re-invoke to continue with interrupt_before @@ -2352,6 +2370,9 @@ async def test_conditional_graph() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2711,6 +2732,9 @@ async def test_conditional_graph_state() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) await app_w_interrupt.aupdate_state( @@ -2752,6 +2776,9 @@ async def test_conditional_graph_state() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2825,6 +2852,9 @@ async def test_conditional_graph_state() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) # test state get/update methods with interrupt_before @@ -2877,6 +2907,9 @@ async def test_conditional_graph_state() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) await app_w_interrupt.aupdate_state( @@ -2918,6 +2951,9 @@ async def test_conditional_graph_state() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2991,6 +3027,9 @@ async def test_conditional_graph_state() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) @@ -3875,6 +3914,9 @@ async def test_state_graph_packets() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) # modify ai message @@ -3927,6 +3969,9 @@ async def test_state_graph_packets() -> None: } }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -4033,6 +4078,9 @@ async def test_state_graph_packets() -> None: }, }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) await app_w_interrupt.aupdate_state( @@ -4078,6 +4126,9 @@ async def test_state_graph_packets() -> None: "step": 5, "writes": {"agent": {"messages": AIMessage(content="answer", id="ai2")}}, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) @@ -4306,6 +4357,9 @@ async def test_message_graph() -> None: ) }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) # modify ai message @@ -4352,6 +4406,9 @@ async def test_message_graph() -> None: ) }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -4423,6 +4480,9 @@ async def test_message_graph() -> None: ) }, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) await app_w_interrupt.aupdate_state( @@ -4464,6 +4524,9 @@ async def test_message_graph() -> None: "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, ) @@ -6016,6 +6079,1136 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 +@pytest.mark.repeat(10) +@pytest.mark.parametrize( + "checkpointer_fct", + [ + lambda: MemorySaverAssertImmutable(put_sleep=0.2), + lambda: AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_nested_graph_interrupts( + checkpointer_fct: Callable[[], BaseCheckpointSaver], +) -> None: + try: + checkpointer = checkpointer_fct() + + class InnerState(TypedDict): + my_key: str + my_other_key: str + + async def inner_1(state: InnerState): + return { + "my_key": state["my_key"] + " here", + "my_other_key": state["my_key"], + } + + async def inner_2(state: InnerState): + return { + "my_key": state["my_key"] + " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + async def outer_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + async def outer_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("outer_1", outer_1) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_2", outer_2) + graph.set_entry_point("outer_1") + graph.add_edge("outer_1", "inner") + graph.add_edge("inner", "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + {"outer_1": {"my_key": "hi my value"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"inner": {"my_key": "hi my value here and there"}}, + {"outer_2": {"my_key": "hi my value here and there and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, stream_mode="values" + ) + ] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + + # test interrupts BEFORE the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, stream_mode="values" + ) + ] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=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")] == [] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # test interrupts AFTER the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, stream_mode="values" + ) + ] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there and back again", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # test restarting from thread_ts + config = {"configurable": {"thread_id": "6"}} + app = graph.compile(checkpointer=checkpointer) + await app.ainvoke({"my_key": "my value"}, config, debug=True) + + state_history = [c async for c in app.aget_state_history(config)] + assert state_history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + child_state_history = [ + c + async for c in app.aget_state_history( + {"configurable": {"thread_id": "6-inner"}} + ) + ] + assert child_state_history == [ + StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + ), + ] + + # check that child snapshot matches id of parent + child_snapshot = child_state_history[0] + assert ( + child_snapshot.config["configurable"]["thread_ts"] + == state_history[0].config["configurable"]["thread_ts"] + ) + # check resuming from interrupt w/ thread_ts + interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] + before_interrupt_config = before_interrupt_state_snapshot.config + # going to get to interrupt again here + assert await app.ainvoke(None, before_interrupt_config, debug=True) == { + "my_key": "hi my value" + } + # one more "identical" snapshot than before, at top of list + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # going to resume from interrupt + interrupt_config = interrupt_state_snapshot.config + assert (await app.ainvoke(None, interrupt_config, debug=True)) == { + "my_key": "hi my value here and there and back again", + } + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_nested_graph_interrupts_parallel( + checkpointer: BaseCheckpointSaver, +) -> None: + try: + + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + async def inner_1(state: InnerState): + await asyncio.sleep(0.1) + return {"my_key": "got here", "my_other_key": state["my_key"]} + + async def inner_2(state: InnerState): + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + async def outer_1(state: State): + return {"my_key": " and parallel"} + + async def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + "my_key": "", + } + + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "got here and there and parallel and back again", + } + + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": ""}, config)] == [ + # we got to parallel node first + {"outer_1": {"my_key": " and parallel"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + { + "my_key": "", + }, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "got here and there and parallel", + }, + { + "my_key": "got here and there and parallel and back again", + }, + ] + + # # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [{"my_key": ""}] + # 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")] == [] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "got here and there and parallel", + }, + { + "my_key": "got here and there and parallel and back again", + }, + ] + + # test interrupts AFTER the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [{"my_key": ""}] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": "got here and there and parallel"}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "got here and there and parallel and back again", + }, + ] + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_doubly_nested_graph_interrupts( + checkpointer: BaseCheckpointSaver, +) -> None: + try: + + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + async def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + async def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node("child_1", grandchild.compile(interrupt_before=["grandchild_2"])) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + async def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + async def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } + + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + {"parent_1": {"my_key": "hi my value"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, stream_mode="values" + ) + ] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, None) + + async def test_checkpoint_metadata() -> None: """This test verifies that a run's configurable fields are merged with the previous checkpoint config for each step in the run.