From f5390a4ee57abb1ae9f4c1978e0ea52bd6adc70b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 18:04:55 -0700 Subject: [PATCH 01/19] Port to new architecture --- .../langgraph/checkpoint/aiosqlite.py | 2 +- libs/langgraph/langgraph/checkpoint/sqlite.py | 2 +- libs/langgraph/langgraph/constants.py | 12 +- libs/langgraph/langgraph/errors.py | 12 + libs/langgraph/langgraph/pregel/__init__.py | 29 +- libs/langgraph/langgraph/pregel/algo.py | 21 +- libs/langgraph/langgraph/pregel/executor.py | 15 +- libs/langgraph/langgraph/pregel/loop.py | 77 +++- libs/langgraph/tests/test_pregel.py | 375 ++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 401 ++++++++++++++++++ 10 files changed, 923 insertions(+), 23 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/aiosqlite.py b/libs/langgraph/langgraph/checkpoint/aiosqlite.py index 431684d1a..d8da20de2 100644 --- a/libs/langgraph/langgraph/checkpoint/aiosqlite.py +++ b/libs/langgraph/langgraph/checkpoint/aiosqlite.py @@ -246,7 +246,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): # find the latest checkpoint for the thread_id if config["configurable"].get("thread_ts"): await cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", + "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts <= ? ORDER BY thread_ts DESC LIMIT 1", ( str(config["configurable"]["thread_id"]), str(config["configurable"]["thread_ts"]), diff --git a/libs/langgraph/langgraph/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index 6ac9ee593..737d426b7 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -245,7 +245,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): # find the latest checkpoint for the thread_id if config["configurable"].get("thread_ts"): cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", + "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts <= ? ORDER BY thread_ts DESC LIMIT 1", ( str(config["configurable"]["thread_id"]), str(config["configurable"]["thread_ts"]), 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..5c288aa50 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,9 @@ def prepare_next_tasks( PregelTaskWrites(name, writes, triggers), config, ), + CONFIG_KEY_CHECKPOINTER: checkpointer, + CONFIG_KEY_RESUMING: is_resuming, + "thread_id": thread_id, }, ), triggers, diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 352b8cf51..5e17123a2 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -6,7 +6,6 @@ from contextvars import copy_context from types import TracebackType from typing import ( AsyncContextManager, - Awaitable, Callable, Iterator, Optional, @@ -18,6 +17,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") @@ -42,6 +43,10 @@ def BackgroundExecutor(config: RunnableConfig) -> Iterator[Submit]: def done(task: concurrent.futures.Future) -> 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 + tasks.pop(task) except BaseException: pass else: @@ -79,7 +84,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): def submit( self, - fn: Callable[P, Awaitable[T]], + fn: Callable[P, T], *args: P.args, __name__: Optional[str] = None, __cancel_on_exit__: bool = False, @@ -97,12 +102,16 @@ 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: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 5d02080b0..ac7a5ac00 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -38,7 +38,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 +67,7 @@ if TYPE_CHECKING: V = TypeVar("V") INPUT_DONE = object() +INPUT_RESUMING = object() class PregelLoop: @@ -95,6 +97,7 @@ class PregelLoop: ] tasks: Sequence[PregelExecutableTask] stream: deque[Tuple[str, Any]] + is_nested: bool # public @@ -133,7 +136,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 +168,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 +190,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 @@ -204,7 +212,10 @@ class PregelLoop: # 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(self) + else: + return False # produce debug output self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks)) @@ -214,8 +225,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, @@ -236,14 +262,9 @@ 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, @@ -313,6 +334,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): # context manager def __enter__(self) -> Self: + self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) @@ -348,6 +370,20 @@ class SyncPregelLoop(PregelLoop, ContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # handle interrupt + if exc_type is GraphInterrupt: + if exc_value.args[0] is self: + # interrupt raised by this loop + exc_value.args = (object(),) + else: + # interrupt raised by a nested loop, save interrupt checkpoint + self._put_checkpoint({"source": "interrupt"}) + if not self.is_nested: + # in outer graph, catch interrupt + del self.graph + return True and self.stack.__exit__(None, None, None) + + # unwind stack del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) @@ -380,6 +416,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): # context manager async def __aenter__(self) -> Self: + self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) saved = ( await self.checkpointer.aget_tuple(self.config) if self.checkpointer @@ -417,6 +454,22 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # handle interrupt + if exc_type is GraphInterrupt: + if exc_value.args[0] is self: + # interrupt raised by this loop + exc_value.args = (object(),) + else: + # interrupt raised by a nested loop, save interrupt checkpoint + self._put_checkpoint({"source": "interrupt"}) + if not self.is_nested: + # in outer graph, catch interrupt + del self.graph + return True and await asyncio.shield( + self.stack.__aexit__(None, None, None) + ) + + # unwind stack del self.graph return await asyncio.shield( self.stack.__aexit__(exc_type, exc_value, traceback) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2646a64fd..89db73b2a 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7524,6 +7524,381 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: + try: + + 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 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)] == [ + {"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", + }, + ] + # 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": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + + # 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", + }, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there 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_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): + 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", + } + + # 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 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0e927e518..a43d22f8a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6016,6 +6016,407 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: + try: + + 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 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)] == [ + {"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", + }, + ] + # 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": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + + # 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 [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there 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_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): + 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", + } + + # 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. From 375d9a45d91f6e75f96c05b759f7f2132c377e63 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 18:15:47 -0700 Subject: [PATCH 02/19] Update loop.py --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index ac7a5ac00..2d498cee1 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -381,7 +381,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): if not self.is_nested: # in outer graph, catch interrupt del self.graph - return True and self.stack.__exit__(None, None, None) + return True or self.stack.__exit__(None, None, None) # unwind stack del self.graph From afe3905958953899d1cb9f1b90e6589a5e2409e7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 18:15:54 -0700 Subject: [PATCH 03/19] Update loop.py --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 2d498cee1..4f817a7c0 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -465,7 +465,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): if not self.is_nested: # in outer graph, catch interrupt del self.graph - return True and await asyncio.shield( + return True or await asyncio.shield( self.stack.__aexit__(None, None, None) ) From d635a22302df57ad8df7fca2aa61c20ace982ea3 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 12:14:33 -0400 Subject: [PATCH 04/19] sort memory checkpoints monotonically decreasing --- libs/langgraph/langgraph/checkpoint/memory.py | 4 +- libs/langgraph/tests/test_pregel.py | 140 +++++++++--------- libs/langgraph/tests/test_pregel_async.py | 140 +++++++++--------- 3 files changed, 143 insertions(+), 141 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index bd5dc6fd1..eec3fcdf7 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -119,7 +119,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) 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 diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 89db73b2a..f6f1257aa 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -608,33 +608,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: # list history assert [c for c in app.get_state_history({"configurable": {"thread_id": 1}})] == [ 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,47 +616,8 @@ 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, ), StateSnapshot( @@ -694,12 +629,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}, + created_at=AnyStr(), parent_config=None, ), 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=None, + ), + 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=None, + ), + 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=None, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, next=(), config={ "configurable": { @@ -707,8 +681,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=None, + ), + 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=None, + ), + 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, ), ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a43d22f8a..7e60c77fb 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -727,33 +727,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N c async for c in app.aget_state_history({"configurable": {"thread_id": 1}}) ] == [ 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,47 +735,8 @@ 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, ), StateSnapshot( @@ -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}, + created_at=AnyStr(), parent_config=None, ), 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=None, + ), + 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=None, + ), + 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=None, + ), + 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=None, + ), + 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=None, + ), + 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, ), ] From f2a95ad67f73099f9eeb26989af501c1d0765dd3 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 12:18:12 -0400 Subject: [PATCH 05/19] fix broken interrupt tests --- libs/langgraph/tests/test_pregel.py | 1 + libs/langgraph/tests/test_pregel_async.py | 1 + 2 files changed, 2 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f6f1257aa..1a2314a48 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7687,6 +7687,7 @@ def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> 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): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7e60c77fb..b168a7f08 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6196,6 +6196,7 @@ async def test_nested_graph_interrupts_parallel( 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): From 77f6ec5a2b7f165d136f66d7b01922a5ad2ce28e Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 13:03:30 -0400 Subject: [PATCH 06/19] update memory --- libs/langgraph/langgraph/checkpoint/memory.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index eec3fcdf7..f695d0f42 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -70,18 +70,20 @@ class MemorySaver(BaseCheckpointSaver): Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ thread_id = config["configurable"]["thread_id"] - if ts := config["configurable"].get("thread_ts"): - if saved := self.storage[thread_id].get(ts): - checkpoint, metadata = saved - writes = self.writes[(thread_id, ts)] - return CheckpointTuple( - config=config, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - pending_writes=[ - (id, c, self.serde.loads(v)) for id, c, v in writes - ], - ) + if thread_ts := config["configurable"].get("thread_ts"): + if checkpoints := self.storage[thread_id]: + ts = max(key for key in checkpoints.keys() if key <= thread_ts) + if saved := self.storage[thread_id].get(ts): + checkpoint, metadata = saved + writes = self.writes[(thread_id, ts)] + return CheckpointTuple( + config=config, + checkpoint=self.serde.loads(checkpoint), + metadata=self.serde.loads(metadata), + pending_writes=[ + (id, c, self.serde.loads(v)) for id, c, v in writes + ], + ) else: if checkpoints := self.storage[thread_id]: ts = max(checkpoints.keys()) From 6e57fa5f9ab5594df0cba815bc3d7c44628c9ab9 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 14:32:57 -0400 Subject: [PATCH 07/19] return future on self._put_checkpoint --- libs/langgraph/langgraph/pregel/loop.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 4f817a7c0..a3e29564d 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 @@ -269,7 +270,7 @@ class PregelLoop: def _put_checkpoint( self, metadata: CheckpointMetadata, - ) -> None: + ) -> concurrent.futures.Future: # assign step metadata["step"] = self.step # bail if no checkpointer @@ -280,7 +281,7 @@ class PregelLoop: self.checkpoint, self.channels, self.step ) # save it, without blocking - self.submit( + fut = self.submit( self.checkpointer_put, self.checkpoint_config, copy_checkpoint(self.checkpoint), @@ -304,8 +305,12 @@ class PregelLoop: self.checkpoint_metadata, ) ) + else: + fut = concurrent.futures.Future() + fut.set_result(None) # increment step self.step += 1 + return fut class SyncPregelLoop(PregelLoop, ContextManager): @@ -377,7 +382,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): exc_value.args = (object(),) else: # interrupt raised by a nested loop, save interrupt checkpoint - self._put_checkpoint({"source": "interrupt"}) + self._put_checkpoint({"source": "interrupt"}).result() if not self.is_nested: # in outer graph, catch interrupt del self.graph From fbced023f364a69d70dd6f8529799e35a4b7f5c4 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 15:06:15 -0400 Subject: [PATCH 08/19] add more tests --- libs/langgraph/langgraph/checkpoint/memory.py | 3 +- libs/langgraph/tests/test_pregel.py | 894 ++++++++++++++++- libs/langgraph/tests/test_pregel_async.py | 912 +++++++++++++++++- libs/langgraph/tests/utils.py | 17 + 4 files changed, 1822 insertions(+), 4 deletions(-) create mode 100644 libs/langgraph/tests/utils.py diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index f695d0f42..2af8fc21d 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -72,7 +72,8 @@ class MemorySaver(BaseCheckpointSaver): thread_id = config["configurable"]["thread_id"] if thread_ts := config["configurable"].get("thread_ts"): if checkpoints := self.storage[thread_id]: - ts = max(key for key in checkpoints.keys() if key <= thread_ts) + matching_keys = [key for key in checkpoints.keys() if key <= thread_ts] + ts = max(matching_keys) if matching_keys else None if saved := self.storage[thread_id].get(ts): checkpoint, metadata = saved writes = self.writes[(thread_id, ts)] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 1a2314a48..e651a7643 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -67,6 +67,7 @@ from tests.memory_assert import ( MemorySaverAssertImmutable, NoopSerializer, ) +from tests.utils import assert_state_history_equal def test_graph_validation() -> None: @@ -7586,10 +7587,161 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: assert app.invoke({"my_key": "my value"}, config, debug=True) == { "my_key": "hi my value", } - + assert_state_history_equal( + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert app.invoke(None, config, debug=True) == { "my_key": "hi my value here and there and back again", } + assert_state_history_equal( + 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": 4, + }, + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} @@ -7631,8 +7783,102 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: "my_key": "hi my value", }, ] + assert_state_history_equal( + 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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) # while we're waiting for the node w/ interrupt inside to finish assert [*app.stream(None, config, stream_mode="values")] == [] + assert_state_history_equal( + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert [*app.stream(None, config, stream_mode="values")] == [ { "my_key": "hi my value here and there", @@ -7641,6 +7887,106 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: "my_key": "hi my value here and there and back again", }, ] + assert_state_history_equal( + 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": 4, + }, + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) # test interrupts AFTER the node w/ interrupts app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) @@ -7653,16 +7999,562 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: "my_key": "hi my value", }, ] + assert_state_history_equal( + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert [*app.stream(None, config, stream_mode="values")] == [ { "my_key": "hi my value here and there", }, ] + assert_state_history_equal( + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert [*app.stream(None, config, stream_mode="values")] == [ { "my_key": "hi my value here and there and back again", }, ] + assert_state_history_equal( + 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": 4, + }, + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + + # test restarting from thread_ts + config = {"configurable": {"thread_id": "6"}} + app = graph.compile(checkpointer=checkpointer) + app.invoke({"my_key": "my value"}, config, debug=True) + state_history = [c for c in app.get_state_history(config)] + assert_state_history_equal( + state_history, + [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "interrupt", "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": "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + child_state_history = [ + c for c in app.get_state_history({"configurable": {"thread_id": "6-inner"}}) + ] + assert_state_history_equal( + 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(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=(), + config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "hi my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + + # check that parent snapshot is always older than child + 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) is None + assert_state_history_equal( + list(app.get_state_history(config)), + [ + # NOTE: there is an identical snapshot here since we replayed from before interrupt + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "interrupt", "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": "interrupt", "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": "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + # 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_state_history_equal( + 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": 4, + }, + 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": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + ), + # NOTE: there is an identical snapshot here since we replayed from before interrupt + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "interrupt", "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": "interrupt", "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": "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) finally: if hasattr(checkpointer, "__exit__"): checkpointer.__exit__(None, None, None) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b168a7f08..ba0876d40 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -61,6 +61,7 @@ from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, ) +from tests.utils import assert_state_history_equal async def test_checkpoint_errors() -> None: @@ -6078,11 +6079,163 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { "my_key": "hi my value", } - + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert await app.ainvoke(None, config, debug=True) == { "my_key": "hi my value here and there and back again", } - + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": 4, + }, + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ @@ -6133,8 +6286,104 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non "my_key": "hi my value", }, ] + await asyncio.sleep(0.05) + assert_state_history_equal( + [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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) # 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")] == [] + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert [c async for c in app.astream(None, config, stream_mode="values")] == [ { "my_key": "hi my value here and there", @@ -6143,6 +6392,107 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non "my_key": "hi my value here and there and back again", }, ] + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": 4, + }, + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) # test interrupts AFTER the node w/ interrupts app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) @@ -6160,16 +6510,574 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non "my_key": "hi my value", }, ] + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert [c async for c in app.astream(None, config, stream_mode="values")] == [ { "my_key": "hi my value here and there", }, ] + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) assert [c async for c in app.astream(None, config, stream_mode="values")] == [ { "my_key": "hi my value here and there and back again", }, ] + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": 4, + }, + 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": 3, + }, + 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": "interrupt", "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + + # 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) + await asyncio.sleep(0.05) + state_history = [c async for c in app.aget_state_history(config)] + ( + state_history, + [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "interrupt", "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": "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, + ), + ], + # ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + + await asyncio.sleep(0.05) + child_state_history = [ + c + async for c in app.aget_state_history( + {"configurable": {"thread_id": "6-inner"}} + ) + ] + assert_state_history_equal( + 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(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=(), + config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "hi my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + + # check that parent snapshot is always older than child + 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 (await app.ainvoke(None, before_interrupt_config, debug=True)) is None + + await asyncio.sleep(0.05) + assert_state_history_equal( + [s async for s in app.aget_state_history(config)], + [ + # NOTE: there is an identical snapshot here since we replayed from before interrupt + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "interrupt", "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": "interrupt", "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": "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) + # going to restart 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", + } + await asyncio.sleep(0.05) + assert_state_history_equal( + [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": 4, + }, + 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": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "thread_ts": AnyStr(), + } + }, + ), + # NOTE: there is an identical snapshot here since we replayed from before interrupt + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + metadata={"source": "interrupt", "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": "interrupt", "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": "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, + ), + ], + ignore_parent_config=isinstance(checkpointer, MemorySaver), + ) finally: if hasattr(checkpointer, "__aexit__"): await checkpointer.__aexit__(None, None, None) diff --git a/libs/langgraph/tests/utils.py b/libs/langgraph/tests/utils.py new file mode 100644 index 000000000..22e46a3ac --- /dev/null +++ b/libs/langgraph/tests/utils.py @@ -0,0 +1,17 @@ +from langgraph.pregel import StateSnapshot + + +def assert_state_history_equal( + actual_state_history: list[StateSnapshot], + expected_state_history: list[StateSnapshot], + ignore_parent_config: bool = False, +) -> None: + assert ( + len(actual_state_history) == len(expected_state_history) + ), f"Got different lengths for state history: {len(actual_state_history)} for actual, {len(expected_state_history)} for expected" + for actual, expected in zip(actual_state_history, expected_state_history): + if ignore_parent_config: + actual = actual._replace(parent_config=None) + expected = expected._replace(parent_config=None) + + assert actual == expected From 0c1ec8d5d202ebe5dfb50a78142f52cd701786c1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 16:30:08 -0700 Subject: [PATCH 09/19] Update sync bg executor to wait on all tasks on exit --- libs/langgraph/langgraph/pregel/executor.py | 94 +++++++++++++-------- 1 file changed, 57 insertions(+), 37 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 5e17123a2..948711508 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -1,13 +1,13 @@ 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, Callable, - Iterator, + ContextManager, Optional, Protocol, TypeVar, @@ -35,45 +35,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: + 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: + 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 + concurrent.futures.wait({t for t in self.tasks if not t.done()}) + # shutdown the executor + self.stack.__exit__(exc_type, exc_value, traceback) + # raise caught exception + if exc_type is not None: + raise exc_value + # re-raise the first exception that occurred in a task + for task in self.tasks: 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 - tasks.pop(task) - except BaseException: + except concurrent.futures.CancelledError: 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 - - 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() class AsyncBackgroundExecutor(AsyncContextManager): @@ -131,6 +149,8 @@ class AsyncBackgroundExecutor(AsyncContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: + # 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 for task, cancel in self.tasks.items(): if cancel: task.cancel(self.sentinel) From 42174ad9fb2fada887803094e1e0cfd70e771832 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 16:30:19 -0700 Subject: [PATCH 10/19] Add optional id arg to create_checkpoint --- libs/langgraph/langgraph/channels/manager.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) 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"], From 98a962661b4f1683fae4ce79fdc0f9e93e52363d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 16:32:15 -0700 Subject: [PATCH 11/19] Instead of saving an additional interrupt checkpoint, make child graphs keep a single checkpoint for each parent checkpoint - while the inner graph makes progress it overwrites the partial progress checkpoints, eventually keeping only one for each outer step - implement parent_config in MemorySaver - fix edge cases in PregelLoop --- .../langgraph/checkpoint/aiosqlite.py | 2 +- libs/langgraph/langgraph/checkpoint/memory.py | 58 +++++++++++++------ libs/langgraph/langgraph/checkpoint/sqlite.py | 2 +- libs/langgraph/langgraph/pregel/algo.py | 1 + libs/langgraph/langgraph/pregel/debug.py | 14 +++-- libs/langgraph/langgraph/pregel/loop.py | 41 +++++++------ libs/langgraph/tests/memory_assert.py | 3 +- 7 files changed, 73 insertions(+), 48 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/aiosqlite.py b/libs/langgraph/langgraph/checkpoint/aiosqlite.py index d8da20de2..431684d1a 100644 --- a/libs/langgraph/langgraph/checkpoint/aiosqlite.py +++ b/libs/langgraph/langgraph/checkpoint/aiosqlite.py @@ -246,7 +246,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): # find the latest checkpoint for the thread_id if config["configurable"].get("thread_ts"): await cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts <= ? ORDER BY thread_ts DESC LIMIT 1", + "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", ( str(config["configurable"]["thread_id"]), str(config["configurable"]["thread_ts"]), diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 2af8fc21d..b71dd0603 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, @@ -70,25 +70,30 @@ class MemorySaver(BaseCheckpointSaver): Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ thread_id = config["configurable"]["thread_id"] - if thread_ts := config["configurable"].get("thread_ts"): - if checkpoints := self.storage[thread_id]: - matching_keys = [key for key in checkpoints.keys() if key <= thread_ts] - ts = max(matching_keys) if matching_keys else None - if saved := self.storage[thread_id].get(ts): - checkpoint, metadata = saved - writes = self.writes[(thread_id, ts)] - return CheckpointTuple( - config=config, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - pending_writes=[ - (id, c, self.serde.loads(v)) for id, c, v in writes - ], - ) + if ts := config["configurable"].get("thread_ts"): + if saved := self.storage[thread_id].get(ts): + checkpoint, metadata, parent_ts = saved + writes = self.writes[(thread_id, ts)] + return CheckpointTuple( + config=config, + checkpoint=self.serde.loads(checkpoint), + metadata=self.serde.loads(metadata), + 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}}, @@ -97,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( @@ -122,7 +135,7 @@ 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 sorted( + 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 @@ -147,6 +160,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( @@ -172,6 +193,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/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index 737d426b7..6ac9ee593 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -245,7 +245,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): # find the latest checkpoint for the thread_id if config["configurable"].get("thread_ts"): cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts <= ? ORDER BY thread_ts DESC LIMIT 1", + "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", ( str(config["configurable"]["thread_id"]), str(config["configurable"]["thread_ts"]), diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 5c288aa50..13ce7e183 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -382,6 +382,7 @@ def prepare_next_tasks( 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/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a3e29564d..d2abf8372 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,5 +1,4 @@ import asyncio -import concurrent.futures from collections import deque from contextlib import AsyncExitStack, ExitStack from types import TracebackType @@ -69,6 +68,7 @@ if TYPE_CHECKING: V = TypeVar("V") INPUT_DONE = object() INPUT_RESUMING = object() +EMPTY_LIST = [] class PregelLoop: @@ -126,9 +126,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_LIST, + interrupt_after: Sequence[str] = EMPTY_LIST, + interrupt_before: Sequence[str] = EMPTY_LIST, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, ) -> bool: """Execute a single iteration of the Pregel loop. @@ -208,7 +208,12 @@ 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): @@ -267,10 +272,7 @@ class PregelLoop: # done with input self.input = INPUT_RESUMING if is_resuming else INPUT_DONE - def _put_checkpoint( - self, - metadata: CheckpointMetadata, - ) -> concurrent.futures.Future: + def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: # assign step metadata["step"] = self.step # bail if no checkpointer @@ -278,10 +280,17 @@ class PregelLoop: # 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 - fut = self.submit( + self.submit( self.checkpointer_put, self.checkpoint_config, copy_checkpoint(self.checkpoint), @@ -305,12 +314,8 @@ class PregelLoop: self.checkpoint_metadata, ) ) - else: - fut = concurrent.futures.Future() - fut.set_result(None) # increment step self.step += 1 - return fut class SyncPregelLoop(PregelLoop, ContextManager): @@ -380,9 +385,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): if exc_value.args[0] is self: # interrupt raised by this loop exc_value.args = (object(),) - else: - # interrupt raised by a nested loop, save interrupt checkpoint - self._put_checkpoint({"source": "interrupt"}).result() if not self.is_nested: # in outer graph, catch interrupt del self.graph @@ -464,9 +466,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): if exc_value.args[0] is self: # interrupt raised by this loop exc_value.args = (object(),) - else: - # interrupt raised by a nested loop, save interrupt checkpoint - self._put_checkpoint({"source": "interrupt"}) if not self.is_nested: # in outer graph, catch interrupt del self.graph diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index db1212ae9..6709a404a 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -85,7 +85,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 +93,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): self.serde.dumps(checkpoint), # merge configurable fields and metadata self.serde.dumps({**configurable, **metadata}), + thread_ts, ) } ) From 3f860617a5cbc450435a78c1214e0f599022e2e3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 16:32:18 -0700 Subject: [PATCH 12/19] Update tests --- libs/langgraph/poetry.lock | 16 +- libs/langgraph/pyproject.toml | 3 +- libs/langgraph/tests/test_pregel.py | 1653 ++++++++++----------- libs/langgraph/tests/test_pregel_async.py | 1649 ++++++++++---------- 4 files changed, 1563 insertions(+), 1758 deletions(-) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 3bee7a407..4425bb6e3 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 = "5fb6190a1b01d0cd351ea9a0023c8c8d6acf4fe831101ab87f9fc41308f74b74" +content-hash = "da02c84e7232f1aa2980537e71d88a757a62444a41b1873566ec81842fda8f82" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 7b4ab2078..fd9392802 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/test_pregel.py b/libs/langgraph/tests/test_pregel.py index e651a7643..ee6545120 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, @@ -67,7 +68,6 @@ from tests.memory_assert import ( MemorySaverAssertImmutable, NoopSerializer, ) -from tests.utils import assert_state_history_equal def test_graph_validation() -> None: @@ -607,7 +607,8 @@ 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={"inbox": 4, "output": 5, "input": 3}, next=(), @@ -619,7 +620,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "loop", "step": 6, "writes": 5}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][1].config, ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, @@ -632,7 +633,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "loop", "step": 5, "writes": None}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][2].config, ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 3}, @@ -645,7 +646,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "input", "step": 4, "writes": 3}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][3].config, ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 20}, @@ -658,7 +659,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "loop", "step": 3, "writes": None}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][4].config, ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 20}, @@ -671,7 +672,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "input", "step": 2, "writes": 20}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][5].config, ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 2}, @@ -684,7 +685,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "loop", "step": 1, "writes": 4}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][6].config, ), StateSnapshot( values={"inbox": 3, "input": 2}, @@ -697,7 +698,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: }, metadata={"source": "loop", "step": 0, "writes": None}, created_at=AnyStr(), - parent_config=None, + parent_config=[*app.checkpointer.list(thread1)][7].config, ), StateSnapshot( values={"input": 2}, @@ -1724,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"][ @@ -1772,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)] == [ @@ -1879,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 @@ -1929,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( @@ -1971,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)] == [ @@ -2078,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 @@ -2128,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)] == [ @@ -2527,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( @@ -2566,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)] == [ @@ -2637,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 @@ -2685,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( @@ -2724,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)] == [ @@ -2795,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 @@ -2819,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)] == [ @@ -2855,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)] == [ @@ -2912,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)] == [ @@ -2970,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)] == [ @@ -3027,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)] == [ @@ -4130,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 @@ -4183,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)] == [ @@ -4287,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( @@ -4338,6 +4360,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: } }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -4621,6 +4644,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -4664,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)] == [ @@ -4747,6 +4772,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -4788,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( @@ -4851,6 +4878,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -4897,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)] == [ @@ -4980,6 +5009,7 @@ def test_message_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -5021,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 @@ -5062,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, ) @@ -5343,6 +5375,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -5386,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)] == [ @@ -5469,6 +5503,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -5510,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( @@ -5573,6 +5609,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) # modify ai message @@ -5619,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)] == [ @@ -5702,6 +5740,7 @@ def test_root_graph( ) }, }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) app_w_interrupt.update_state( @@ -5743,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 @@ -5784,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 @@ -5857,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 @@ -6886,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)] == [ @@ -7525,19 +7568,23 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.repeat(10) @pytest.mark.parametrize( - "checkpointer", + "checkpointer_fct", [ - MemorySaverAssertImmutable(), - SqliteSaver.from_conn_string(":memory:"), + lambda: MemorySaverAssertImmutable(), + lambda: SqliteSaver.from_conn_string(":memory:"), ], ids=[ "memory", "sqlite", ], ) -def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: +def test_nested_graph_interrupts( + checkpointer_fct: Callable[[], BaseCheckpointSaver], +) -> None: try: + checkpointer = checkpointer_fct() class InnerState(TypedDict): my_key: str @@ -7587,161 +7634,133 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: assert app.invoke({"my_key": "my value"}, config, debug=True) == { "my_key": "hi my value", } - assert_state_history_equal( - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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_state_history_equal( - 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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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"}} @@ -7783,102 +7802,84 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: "my_key": "hi my value", }, ] - assert_state_history_equal( - 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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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_state_history_equal( - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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", @@ -7887,106 +7888,92 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: "my_key": "hi my value here and there and back again", }, ] - assert_state_history_equal( - 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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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"]) @@ -7999,562 +7986,453 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: "my_key": "hi my value", }, ] - assert_state_history_equal( - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + # 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_state_history_equal( - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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_state_history_equal( - 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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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) - app.invoke({"my_key": "my value"}, config, debug=True) + 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_equal( - state_history, - [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, - metadata={"source": "interrupt", "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": "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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_state_history_equal( - child_state_history, - [ - StateSnapshot( - values={"my_key": "hi my value here"}, - next=(), - config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), + 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", } }, - 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(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=(), - config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "hi my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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 parent snapshot is always older than child + # 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"] + == 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) is None - assert_state_history_equal( - list(app.get_state_history(config)), - [ - # NOTE: there is an identical snapshot here since we replayed from before interrupt - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, - metadata={"source": "interrupt", "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": "interrupt", "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": "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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_state_history_equal( - 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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "thread_ts": AnyStr(), - } - }, - ), - # NOTE: there is an identical snapshot here since we replayed from before interrupt - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, - metadata={"source": "interrupt", "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": "interrupt", "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": "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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) @@ -8844,7 +8722,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 @@ -8919,7 +8797,7 @@ def test_checkpoint_metadata() -> None: # assertions # invoke graph w/o interrupt - app.invoke( + assert app.invoke( {"messages": ["what is weather in sf"]}, { "configurable": { @@ -8928,7 +8806,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 ba0876d40..d520c4092 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, @@ -61,7 +62,6 @@ from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, ) -from tests.utils import assert_state_history_equal async def test_checkpoint_errors() -> None: @@ -724,9 +724,8 @@ 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={"inbox": 4, "output": 5, "input": 3}, next=(), @@ -738,7 +737,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "loop", "step": 6, "writes": 5}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][1].config, ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, @@ -751,7 +750,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "loop", "step": 5, "writes": None}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][2].config, ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 3}, @@ -764,7 +763,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "input", "step": 4, "writes": 3}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][3].config, ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 20}, @@ -777,7 +776,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "loop", "step": 3, "writes": None}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][4].config, ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 20}, @@ -790,7 +789,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "input", "step": 2, "writes": 20}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][5].config, ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 2}, @@ -803,7 +802,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "loop", "step": 1, "writes": 4}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][6].config, ), StateSnapshot( values={"inbox": 3, "input": 2}, @@ -816,7 +815,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N }, metadata={"source": "loop", "step": 0, "writes": None}, created_at=AnyStr(), - parent_config=None, + parent_config=[c async for c in app.checkpointer.alist(thread1)][7].config, ), StateSnapshot( values={"input": 2}, @@ -1937,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( @@ -1981,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)] == [ @@ -2090,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 @@ -2145,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( @@ -2189,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)] == [ @@ -2298,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 @@ -2353,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)] == [ @@ -2712,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( @@ -2753,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)] == [ @@ -2826,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 @@ -2878,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( @@ -2919,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)] == [ @@ -2992,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, ) @@ -3876,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 @@ -3928,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)] == [ @@ -4034,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( @@ -4079,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, ) @@ -4307,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 @@ -4353,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)] == [ @@ -4424,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( @@ -4465,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, ) @@ -6017,19 +6079,23 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 +@pytest.mark.repeat(10) @pytest.mark.parametrize( - "checkpointer", + "checkpointer_fct", [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), + lambda: MemorySaverAssertImmutable(), + lambda: AsyncSqliteSaver.from_conn_string(":memory:"), ], ids=[ "memory", "sqlite", ], ) -async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: +async def test_nested_graph_interrupts( + checkpointer_fct: Callable[[], BaseCheckpointSaver], +) -> None: try: + checkpointer = checkpointer_fct() class InnerState(TypedDict): my_key: str @@ -6080,162 +6146,134 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non "my_key": "hi my value", } await asyncio.sleep(0.05) - assert_state_history_equal( - [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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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", } await asyncio.sleep(0.05) - assert_state_history_equal( - [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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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)] == [ @@ -6287,103 +6325,85 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non }, ] await asyncio.sleep(0.05) - assert_state_history_equal( - [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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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")] == [] await asyncio.sleep(0.05) - assert_state_history_equal( - [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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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", @@ -6393,106 +6413,92 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non }, ] await asyncio.sleep(0.05) - assert_state_history_equal( - [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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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"]) @@ -6511,303 +6517,247 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non }, ] await asyncio.sleep(0.05) - assert_state_history_equal( - [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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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", }, ] await asyncio.sleep(0.05) - assert_state_history_equal( - [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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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", }, ] await asyncio.sleep(0.05) - assert_state_history_equal( - [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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - 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": "interrupt", "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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) - await asyncio.sleep(0.05) + state_history = [c async for c in app.aget_state_history(config)] - ( - state_history, - [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, - metadata={"source": "interrupt", "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": "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, - ), - ], - # ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + 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, + ), + ] await asyncio.sleep(0.05) child_state_history = [ @@ -6816,268 +6766,207 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non {"configurable": {"thread_id": "6-inner"}} ) ] - assert_state_history_equal( - child_state_history, - [ - StateSnapshot( - values={"my_key": "hi my value here"}, - next=(), - config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), + 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", } }, - 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(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=(), - config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "hi my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6-inner", + "thread_ts": AnyStr(), + } + }, + ), + ] - # check that parent snapshot is always older than child + # 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"] + == 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 (await app.ainvoke(None, before_interrupt_config, debug=True)) is None - - await asyncio.sleep(0.05) - assert_state_history_equal( - [s async for s in app.aget_state_history(config)], - [ - # NOTE: there is an identical snapshot here since we replayed from before interrupt - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, - metadata={"source": "interrupt", "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": "interrupt", "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": "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) - # going to restart from interrupt + # 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", } - await asyncio.sleep(0.05) - assert_state_history_equal( - [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(), + 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" } }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 4, - }, - 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": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "thread_ts": AnyStr(), - } - }, - ), - # NOTE: there is an identical snapshot here since we replayed from before interrupt - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, - metadata={"source": "interrupt", "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": "interrupt", "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": "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, - ), - ], - ignore_parent_config=isinstance(checkpointer, MemorySaver), - ) + "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) From aed313ebb576f85312fc96c35309665649ab4300 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 17:09:47 -0700 Subject: [PATCH 13/19] Fix stack not being unwound when suppressing interrupt --- libs/langgraph/langgraph/pregel/executor.py | 55 +++++++----- libs/langgraph/langgraph/pregel/loop.py | 98 ++++++++++----------- 2 files changed, 79 insertions(+), 74 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 948711508..69d3b5d31 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -80,18 +80,18 @@ class BackgroundExecutor(ContextManager): if cancel: task.cancel() # wait for all tasks to finish - concurrent.futures.wait({t for t in self.tasks if not t.done()}) + 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) - # raise caught exception - if exc_type is not None: - raise exc_value # re-raise the first exception that occurred in a task - for task in self.tasks: - try: - task.result() - except concurrent.futures.CancelledError: - pass + 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): @@ -132,16 +132,27 @@ class AsyncBackgroundExecutor(AsyncContextManager): 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, @@ -151,8 +162,6 @@ class AsyncBackgroundExecutor(AsyncContextManager): ) -> Optional[bool]: # 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 - for task, cancel in self.tasks.items(): - if cancel: - task.cancel(self.sentinel) + # 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 da7537cba..bfb8020b0 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -102,6 +102,23 @@ class PregelLoop: # 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 @@ -219,7 +236,7 @@ class PregelLoop: if should_interrupt(self.checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" if self.is_nested: - raise GraphInterrupt(self) + raise GraphInterrupt() else: return False @@ -318,6 +335,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__( @@ -328,24 +354,21 @@ 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 + self.checkpointer_put = checkpointer.put + else: + self.checkpointer_get_next_version = increment + self.checkpointer_put_writes = None + self.checkpointer_put = None # context manager def __enter__(self) -> Self: - self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) @@ -381,16 +404,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - # handle interrupt - if exc_type is GraphInterrupt: - if exc_value.args[0] is self: - # interrupt raised by this loop - exc_value.args = (object(),) - if not self.is_nested: - # in outer graph, catch interrupt - del self.graph - return True or self.stack.__exit__(None, None, None) - # unwind stack del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) @@ -405,26 +418,21 @@ 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 + self.checkpointer_put = checkpointer.aput + else: + self.checkpointer_get_next_version = increment + self.checkpointer_put_writes = None + self.checkpointer_put = None # context manager async def __aenter__(self) -> Self: - self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) saved = ( await self.checkpointer.aget_tuple(self.config) if self.checkpointer @@ -462,18 +470,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - # handle interrupt - if exc_type is GraphInterrupt: - if exc_value.args[0] is self: - # interrupt raised by this loop - exc_value.args = (object(),) - if not self.is_nested: - # in outer graph, catch interrupt - del self.graph - return True or await asyncio.shield( - self.stack.__aexit__(None, None, None) - ) - # unwind stack del self.graph return await asyncio.shield( From cd4fae601cbfcdaaaa5a0b68d64abcf58610c2d0 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 20:26:12 -0400 Subject: [PATCH 14/19] remove sleep --- libs/langgraph/tests/test_pregel.py | 1 - libs/langgraph/tests/test_pregel_async.py | 10 ---------- 2 files changed, 11 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ee6545120..b0a822209 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8457,7 +8457,6 @@ def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> 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): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d520c4092..d5cdfc472 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6145,7 +6145,6 @@ async def test_nested_graph_interrupts( assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { "my_key": "hi my value", } - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, @@ -6187,7 +6186,6 @@ async def test_nested_graph_interrupts( assert await app.ainvoke(None, config, debug=True) == { "my_key": "hi my value here and there and back again", } - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, @@ -6324,7 +6322,6 @@ async def test_nested_graph_interrupts( "my_key": "hi my value", }, ] - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, @@ -6365,7 +6362,6 @@ async def test_nested_graph_interrupts( ] # 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")] == [] - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, @@ -6412,7 +6408,6 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there and back again", }, ] - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, @@ -6516,7 +6511,6 @@ async def test_nested_graph_interrupts( "my_key": "hi my value", }, ] - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, @@ -6560,7 +6554,6 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there", }, ] - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there"}, @@ -6626,7 +6619,6 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there and back again", }, ] - await asyncio.sleep(0.05) assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, @@ -6759,7 +6751,6 @@ async def test_nested_graph_interrupts( ), ] - await asyncio.sleep(0.05) child_state_history = [ c async for c in app.aget_state_history( @@ -6993,7 +6984,6 @@ async def test_nested_graph_interrupts_parallel( 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): From ed9c015d24a281ec9acc6d855972ca12f9c58a05 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 20:43:28 -0400 Subject: [PATCH 15/19] remove utils --- libs/langgraph/tests/utils.py | 17 ----------------- 1 file changed, 17 deletions(-) delete mode 100644 libs/langgraph/tests/utils.py diff --git a/libs/langgraph/tests/utils.py b/libs/langgraph/tests/utils.py deleted file mode 100644 index 22e46a3ac..000000000 --- a/libs/langgraph/tests/utils.py +++ /dev/null @@ -1,17 +0,0 @@ -from langgraph.pregel import StateSnapshot - - -def assert_state_history_equal( - actual_state_history: list[StateSnapshot], - expected_state_history: list[StateSnapshot], - ignore_parent_config: bool = False, -) -> None: - assert ( - len(actual_state_history) == len(expected_state_history) - ), f"Got different lengths for state history: {len(actual_state_history)} for actual, {len(expected_state_history)} for expected" - for actual, expected in zip(actual_state_history, expected_state_history): - if ignore_parent_config: - actual = actual._replace(parent_config=None) - expected = expected._replace(parent_config=None) - - assert actual == expected From fd8f9e87dcfaaf6177ba793a579fe9b77a00c95e Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 20:48:21 -0400 Subject: [PATCH 16/19] update types --- libs/langgraph/langgraph/pregel/executor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 69d3b5d31..5da78a26c 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -6,6 +6,7 @@ from contextvars import copy_context from types import TracebackType from typing import ( AsyncContextManager, + Awaitable, Callable, ContextManager, Optional, @@ -102,7 +103,7 @@ class AsyncBackgroundExecutor(AsyncContextManager): def submit( self, - fn: Callable[P, T], + fn: Callable[P, Awaitable[T]], *args: P.args, __name__: Optional[str] = None, __cancel_on_exit__: bool = False, From ee7cab664ad1934283121e95bdd29f501d47e6e7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 17:52:02 -0700 Subject: [PATCH 17/19] Fix checkpoint put ordering issue --- libs/langgraph/langgraph/pregel/loop.py | 62 ++++++++++++++++++----- libs/langgraph/tests/memory_assert.py | 6 +++ libs/langgraph/tests/test_pregel.py | 2 +- libs/langgraph/tests/test_pregel_async.py | 2 +- 4 files changed, 57 insertions(+), 15 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index bfb8020b0..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 @@ -68,7 +69,7 @@ if TYPE_CHECKING: V = TypeVar("V") INPUT_DONE = object() INPUT_RESUMING = object() -EMPTY_LIST = [] +EMPTY_SEQ = () class PregelLoop: @@ -79,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" @@ -143,9 +152,9 @@ class PregelLoop: def tick( self, *, - output_keys: Union[str, Sequence[str]] = EMPTY_LIST, - interrupt_after: Sequence[str] = EMPTY_LIST, - interrupt_before: Sequence[str] = EMPTY_LIST, + 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. @@ -294,7 +303,7 @@ class PregelLoop: # 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( @@ -308,8 +317,11 @@ class PregelLoop: 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, @@ -360,11 +372,23 @@ class SyncPregelLoop(PregelLoop, ContextManager): if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.put_writes - self.checkpointer_put = checkpointer.put else: self.checkpointer_get_next_version = increment + self._checkpointer_put_after_previous = None self.checkpointer_put_writes = None - self.checkpointer_put = 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 @@ -424,11 +448,23 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.aput_writes - self.checkpointer_put = checkpointer.aput else: self.checkpointer_get_next_version = increment + self._checkpointer_put_after_previous = None self.checkpointer_put_writes = None - self.checkpointer_put = 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 diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index 6709a404a..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): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b0a822209..d90a83ab8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7572,7 +7572,7 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.parametrize( "checkpointer_fct", [ - lambda: MemorySaverAssertImmutable(), + lambda: MemorySaverAssertImmutable(put_sleep=0.2), lambda: SqliteSaver.from_conn_string(":memory:"), ], ids=[ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d5cdfc472..424d63af8 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6083,7 +6083,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.parametrize( "checkpointer_fct", [ - lambda: MemorySaverAssertImmutable(), + lambda: MemorySaverAssertImmutable(put_sleep=0.2), lambda: AsyncSqliteSaver.from_conn_string(":memory:"), ], ids=[ From 0891cc1864c74bfe6c2c19b0937397d136796eff Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 21:18:42 -0400 Subject: [PATCH 18/19] comments / bring back sleep --- libs/langgraph/tests/test_pregel.py | 4 ++++ libs/langgraph/tests/test_pregel_async.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d90a83ab8..23727aded 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8457,6 +8457,7 @@ def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> 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): @@ -8503,6 +8504,9 @@ def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> "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 isnt 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)] == [ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 424d63af8..e906ecf87 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6984,6 +6984,7 @@ async def test_nested_graph_interrupts_parallel( 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): @@ -7030,6 +7031,9 @@ async def test_nested_graph_interrupts_parallel( "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 isnt 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)] == [ From c268d695f0281893b176e250487bd7592d2ec1b1 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 21:28:11 -0400 Subject: [PATCH 19/19] spellcheck --- libs/langgraph/tests/test_pregel.py | 2 +- libs/langgraph/tests/test_pregel_async.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 23727aded..db902eca0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8506,7 +8506,7 @@ def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> # 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 isnt called again (because we dont see outer_1 output again in 2nd stream) + # - 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)] == [ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e906ecf87..c3f8bb3bf 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7033,7 +7033,7 @@ async def test_nested_graph_interrupts_parallel( # 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 isnt called again (because we dont see outer_1 output again in 2nd stream) + # - 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)] == [