From 712d52ad3b3fc557f8828622266afc6a88e906f0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 7 Aug 2024 11:39:14 -0700 Subject: [PATCH 1/3] fix test setup for checkpointers --- libs/langgraph/tests/conftest.py | 183 +- libs/langgraph/tests/test_pregel.py | 3502 ++++++++++--------- libs/langgraph/tests/test_pregel_async.py | 3821 ++++++++++----------- 3 files changed, 3756 insertions(+), 3750 deletions(-) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index eb8f20ce9..f0c85f9a6 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,16 +1,20 @@ -from uuid import UUID +import asyncio +from concurrent.futures import ThreadPoolExecutor +from contextlib import asynccontextmanager, contextmanager +from typing import AsyncIterator, Iterator, TypeVar +from uuid import UUID, uuid4 import pytest -from psycopg import AsyncConnection -from psycopg.errors import UndefinedTable -from psycopg.rows import dict_row +from psycopg import AsyncConnection, Connection from pytest_mock import MockerFixture from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.checkpoint.sqlite import SqliteSaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from tests.memory_assert import MemorySaverAssertImmutable -DEFAULT_POSTGRES_URI = ( - "postgres://postgres:postgres@localhost:5442/postgres?sslmode=disable" -) +DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" @pytest.fixture() @@ -21,29 +25,158 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: return mocker.patch("uuid.uuid4", side_effect=side_effect) +""" +pytest-asyncio doesn't support calling async fixtures with getfixturevalue +so we need to use ThreadPoolExecutor to run the async fixture in a thread +https://github.com/pytest-dev/pytest-asyncio/issues/112#issuecomment-462062890 +""" +T = TypeVar("T") + + +def close_loop(loop: asyncio.AbstractEventLoop) -> None: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + asyncio.set_event_loop(None) + loop.close() + + +@contextmanager +def agen_to_gen(agen: AsyncIterator[T]) -> Iterator[T]: + with ThreadPoolExecutor(1) as bg: + loop = asyncio.new_event_loop() + bg.submit(asyncio.set_event_loop, loop).result() + try: + yield bg.submit(loop.run_until_complete, agen.__aenter__()).result() + finally: + bg.submit( + loop.run_until_complete, agen.__aexit__(None, None, None) + ).result() + bg.submit(close_loop, loop).result() + + +# checkpointer fixtures + + @pytest.fixture(scope="function") -async def conn(): - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - yield conn +def checkpointer_memory(): + yield MemorySaverAssertImmutable() -@pytest.fixture(scope="session", autouse=True) -def setup_before_all_tests(): - with PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI) as checkpointer: - checkpointer.setup() +@pytest.fixture(scope="function") +def checkpointer_sqlite(): + with SqliteSaver.from_conn_string(":memory:") as checkpointer: + yield checkpointer -@pytest.fixture(scope="function", autouse=True) -async def clear_test_db(conn): - """Delete all tables before each test.""" +@pytest.fixture(scope="function") +def checkpointer_sqlite_aio(): + with agen_to_gen(_checkpointer_sqlite_aio()) as checkpointer: + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_sqlite_aio(): + async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: + yield checkpointer + + +@pytest.fixture(scope="function") +def checkpointer_postgres(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") try: - await conn.execute("DELETE FROM checkpoints") - await conn.execute("DELETE FROM checkpoint_blobs") - await conn.execute("DELETE FROM checkpoint_writes") - except UndefinedTable: - pass + # yield checkpointer + with PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + checkpointer.setup() + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") -pytest.register_assert_rewrite("tests.memory_assert") +@pytest.fixture(scope="function") +def checkpointer_postgres_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + with PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + checkpointer.setup() + # setup can't run inside pipeline because of implicit transaction + with checkpointer.conn.pipeline() as pipe: + checkpointer.pipe = pipe + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@pytest.fixture(scope="function") +def checkpointer_postgres_aio(): + with agen_to_gen(_checkpointer_postgres_aio()) as checkpointer: + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_postgres_aio(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + async with AsyncPostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + await checkpointer.setup() + yield checkpointer + finally: + # drop unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@pytest.fixture(scope="function") +def checkpointer_postgres_aio_pipe(): + with agen_to_gen(_checkpointer_postgres_aio_pipe()) as checkpointer: + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_postgres_aio_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + async with AsyncPostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + await checkpointer.setup() + # setup can't run inside pipeline because of implicit transaction + async with checkpointer.conn.pipeline() as pipe: + checkpointer.pipe = pipe + yield checkpointer + finally: + # drop unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4ab248098..684c2df92 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8,7 +8,6 @@ from contextlib import contextmanager from typing import ( Annotated, Any, - Callable, Dict, Generator, List, @@ -41,13 +40,11 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( - BaseCheckpointSaver, Checkpoint, CheckpointMetadata, CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.sqlite import SqliteSaver @@ -64,7 +61,6 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr -from tests.conftest import DEFAULT_POSTGRES_URI from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, @@ -573,386 +569,366 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - SqliteSaver.from_conn_string(":memory:"), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], ) def test_invoke_two_processes_in_out_interrupt( - checkpointer: BaseCheckpointSaver, mocker: MockerFixture + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: - with checkpointer as checkpointer: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert app.invoke(2, thread1) is None + + # inbox == 3 + checkpoint = checkpointer.get(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert app.invoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert app.invoke(20, thread1) is None + + # inbox == 21 + checkpoint = checkpointer.get(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert app.invoke(3, thread1) is None + assert app.invoke(None, thread1) == 5 + + # start execution again, stopping at inbox + assert app.invoke(20, thread2) is None + + # inbox == 21 + snapshot = app.get_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + app.update_state(thread2, 25, as_node="one") + assert app.invoke(None, thread2) == 26 + + # no pending tasks + snapshot = app.get_state(thread2) + assert snapshot.next == () + + # list history + history = [c for c in app.get_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - interrupt_after_nodes=["one"], - ) - thread1 = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} + metadata={"source": "loop", "step": 6, "writes": 5}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": None}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[7].config, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, + created_at=AnyStr(), + parent_config=None, + ), + ] - # start execution, stop at inbox - assert app.invoke(2, thread1) is None + # forking from any previous checkpoint w/out forking should do nothing + assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == [] + assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == [] + assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == [] - # inbox == 3 - checkpoint = checkpointer.get(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = app.update_state(history[0].config, None) + assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [] - # resume execution, finish - assert app.invoke(None, thread1) == 4 + fork_config = app.update_state(history[1].config, None) + assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ + {"two": {"output": 5}} + ] - # start execution again, stop at inbox - assert app.invoke(20, thread1) is None - - # inbox == 21 - checkpoint = checkpointer.get(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 - - # send a new value in, interrupting the previous execution - assert app.invoke(3, thread1) is None - assert app.invoke(None, thread1) == 5 - - # start execution again, stopping at inbox - assert app.invoke(20, thread2) is None - - # inbox == 21 - snapshot = app.get_state(thread2) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) - - # update the state, resume - app.update_state(thread2, 25, as_node="one") - assert app.invoke(None, thread2) == 26 - - # no pending tasks - snapshot = app.get_state(thread2) - assert snapshot.next == () - - # list history - history = [c for c in app.get_state_history(thread1)] - assert history == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 6, "writes": 5}, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 5, "writes": None}, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": 4, "writes": 3}, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 3, "writes": None}, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": 2, "writes": 20}, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 1, "writes": 4}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=history[7].config, - ), - StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": -1, "writes": 2}, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # forking from any previous checkpoint w/out forking should do nothing - assert [ - c for c in app.stream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c for c in app.stream(None, history[1].config, stream_mode="updates") - ] == [] - assert [ - c for c in app.stream(None, history[2].config, stream_mode="updates") - ] == [] - - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = app.update_state(history[0].config, None) - assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [] - - fork_config = app.update_state(history[1].config, None) - assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ - {"two": {"output": 5}} - ] - - fork_config = app.update_state(history[2].config, None) - assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ - {"one": {"inbox": 4}} - ] + fork_config = app.update_state(history[2].config, None) + assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ + {"one": {"inbox": 4}} + ] @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - SqliteSaver.from_conn_string(":memory:"), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], ) def test_fork_always_re_runs_nodes( - checkpointer: BaseCheckpointSaver, mocker: MockerFixture + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: - with checkpointer as checkpointer: - add_one = mocker.Mock(side_effect=lambda _: 1) + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + add_one = mocker.Mock(side_effect=lambda _: 1) - builder = StateGraph(Annotated[int, operator.add]) - builder.add_node("add_one", add_one) - builder.add_edge(START, "add_one") - builder.add_conditional_edges( - "add_one", lambda cnt: "add_one" if cnt < 6 else END - ) - graph = builder.compile(checkpointer=checkpointer) + builder = StateGraph(Annotated[int, operator.add]) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END) + graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} + thread1 = {"configurable": {"thread_id": "1"}} - # start execution, stop at inbox - assert [*graph.stream(1, thread1, stream_mode=["values", "updates"])] == [ - ("values", 1), - ("updates", {"add_one": 1}), - ("values", 2), - ("updates", {"add_one": 1}), - ("values", 3), - ("updates", {"add_one": 1}), - ("values", 4), - ("updates", {"add_one": 1}), - ("values", 5), - ("updates", {"add_one": 1}), - ("values", 6), - ] + # start execution, stop at inbox + assert [*graph.stream(1, thread1, stream_mode=["values", "updates"])] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] - # list history - history = [c for c in graph.get_state_history(thread1)] - assert history == [ - StateSnapshot( - values=6, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values=5, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values=4, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values=3, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values=2, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values=1, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values=0, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": -1, "writes": 1}, - created_at=AnyStr(), - parent_config=None, - ), - ] + # list history + history = [c for c in graph.get_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values=5, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values=4, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values=3, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values=2, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values=1, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values=0, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 1}, + created_at=AnyStr(), + parent_config=None, + ), + ] - # forking from any previous checkpoint w/out forking should do nothing - assert [ - c for c in graph.stream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c for c in graph.stream(None, history[1].config, stream_mode="updates") - ] == [] + # forking from any previous checkpoint w/out forking should do nothing + assert [ + c for c in graph.stream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c for c in graph.stream(None, history[1].config, stream_mode="updates") + ] == [] - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = graph.update_state(history[0].config, None) - assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [] + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = graph.update_state(history[0].config, None) + assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [] - fork_config = graph.update_state(history[1].config, None) - assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ - {"add_one": 1} - ] + fork_config = graph.update_state(history[1].config, None) + assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ + {"add_one": 1} + ] - fork_config = graph.update_state(history[2].config, None) - assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ - {"add_one": 1}, - {"add_one": 1}, - ] + fork_config = graph.update_state(history[2].config, None) + assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ + {"add_one": 1}, + {"add_one": 1}, + ] def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: @@ -1277,88 +1253,84 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - SqliteSaver.from_conn_string(":memory:"), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], ) -def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None: - with checkpointer as checkpointer: +def test_pending_writes_resume( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): - value: Annotated[int, operator.add] + class State(TypedDict): + value: Annotated[int, operator.add] - class AwhileMaker: - def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None: - self.sleep = sleep - self.rtn = rtn - self.reset() + class AwhileMaker: + def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None: + self.sleep = sleep + self.rtn = rtn + self.reset() - def __call__(self, input: State) -> Any: - self.calls += 1 - time.sleep(self.sleep) - if isinstance(self.rtn, Exception): - raise self.rtn - else: - return self.rtn + def __call__(self, input: State) -> Any: + self.calls += 1 + time.sleep(self.sleep) + if isinstance(self.rtn, Exception): + raise self.rtn + else: + return self.rtn - def reset(self): - self.calls = 0 + def reset(self): + self.calls = 0 - one = AwhileMaker(0.2, {"value": 2}) - two = AwhileMaker(0.6, ConnectionError("I'm not good")) - builder = StateGraph(State) - builder.add_node("one", one) - builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) - builder.add_edge(START, "one") - builder.add_edge(START, "two") - graph = builder.compile(checkpointer=checkpointer) + one = AwhileMaker(0.2, {"value": 2}) + two = AwhileMaker(0.6, ConnectionError("I'm not good")) + builder = StateGraph(State) + builder.add_node("one", one) + builder.add_node("two", two, retry=RetryPolicy(max_attempts=2)) + builder.add_edge(START, "one") + builder.add_edge(START, "two") + graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - with pytest.raises(ConnectionError, match="I'm not good"): - graph.invoke({"value": 1}, thread1) + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ConnectionError, match="I'm not good"): + graph.invoke({"value": 1}, thread1) - # both nodes should have been called once - assert one.calls == 1 - assert two.calls == 2 # two attempts + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 2 # two attempts - # latest checkpoint should be before nodes "one", "two" - state = graph.get_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") - assert state.metadata == {"source": "loop", "step": 0, "writes": None} - # should contain pending write of "one" - checkpoint = checkpointer.get_tuple(thread1) - assert checkpoint is not None - assert checkpoint.pending_writes == [ - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - ] - # both pending writes come from same task - assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0] + # latest checkpoint should be before nodes "one", "two" + state = graph.get_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + assert state.metadata == {"source": "loop", "step": 0, "writes": None} + # should contain pending write of "one" + checkpoint = checkpointer.get_tuple(thread1) + assert checkpoint is not None + assert checkpoint.pending_writes == [ + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + ] + # both pending writes come from same task + assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0] - # resume execution - with pytest.raises(ConnectionError, match="I'm not good"): - graph.invoke(None, thread1) + # resume execution + with pytest.raises(ConnectionError, match="I'm not good"): + graph.invoke(None, thread1) - # node "one" succeeded previously, so shouldn't be called again - assert one.calls == 1 - # node "two" should have been called once again - assert two.calls == 4 # two attempts before + two attempts now + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 4 # two attempts before + two attempts now - # confirm no new checkpoints saved - state_two = graph.get_state(thread1) - assert state_two == state + # confirm no new checkpoints saved + state_two = graph.get_state(thread1) + assert state_two == state - # resume execution, without exception - two.rtn = {"value": 3} - # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert graph.invoke(None, thread1) == {"value": 6} + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert graph.invoke(None, thread1) == {"value": 6} def test_cond_edge_after_send() -> None: @@ -7577,1422 +7549,1402 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.repeat(10) @pytest.mark.parametrize( - "checkpointer_fct", - [ - lambda: MemorySaverAssertImmutable(put_sleep=0.2), - lambda: SqliteSaver.from_conn_string(":memory:"), - lambda: PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - lambda: PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], ) def test_nested_graph_interrupts( - checkpointer_fct: Callable[[], BaseCheckpointSaver], + request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - with checkpointer_fct() as checkpointer: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class InnerState(TypedDict): - my_key: str - my_other_key: str + 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_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"], - } + 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") + 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 + class State(TypedDict): + my_key: str - def outer_1(state: State): - return {"my_key": "hi " + state["my_key"]} + 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"} + 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") + 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) + 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) == { + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert app.invoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [*app.stream({"my_key": "my value"}, config)] == [ + {"outer_1": {"my_key": "hi my value"}}, + ] + assert [*app.stream(None, config)] == [ + {"inner": {"my_key": "hi my value here and there"}}, + {"outer_2": {"my_key": "hi my value here and there and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { "my_key": "hi my value", - } - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert app.invoke(None, config, debug=True) == { + }, + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { "my_key": "hi my value here and there and back again", - } - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + }, + ] - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [*app.stream({"my_key": "my value"}, config)] == [ - {"outer_1": {"my_key": "hi my value"}}, - ] - assert [*app.stream(None, config)] == [ - {"inner": {"my_key": "hi my value here and there"}}, - {"outer_2": {"my_key": "hi my value here and there and back again"}}, - ] + # test interrupts BEFORE the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) + config = {"configurable": {"thread_id": "4"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [*app.stream(None, config, stream_mode="values")] == [] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # 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", + # test interrupts AFTER the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) + config = {"configurable": {"thread_id": "5"}} + assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + # interrupted after "inner" + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value", + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, }, - ] - assert [*app.stream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value here and there and back again", + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - ] + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [*app.stream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there and back again", + }, + ] + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # 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", + # test restarting from checkpoint_id + config = {"configurable": {"thread_id": "6"}} + app = graph.compile(checkpointer=checkpointer) + assert app.invoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value" + } + state_history = [c for c in app.get_state_history(config)] + assert state_history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value", + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, }, - ] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [*app.stream(None, config, stream_mode="values")] == [] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": 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", + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value here and there and back again", + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test interrupts AFTER the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) - config = {"configurable": {"thread_id": "5"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ - { - "my_key": "my value", + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value", + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - ] - # interrupted after "inner" - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": 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", + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, }, - ] - 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", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [*app.stream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there and back again", - }, - ] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test restarting from checkpoint_id - config = {"configurable": {"thread_id": "6"}} - app = graph.compile(checkpointer=checkpointer) - assert app.invoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value" - } - state_history = [c for c in app.get_state_history(config)] - assert state_history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": 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", "checkpoint_ns": "inner"}} - ) - ] - assert child_state_history == [ - StateSnapshot( - values={"my_key": "hi my value here"}, - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": 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", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - # there should be a single child checkpoint because we only keep - # one child checkpoint per parent checkpoint (in which child ran) - ] - - # check that child snapshot matches id of parent - child_snapshot = child_state_history[0] - assert ( - child_snapshot.config["configurable"]["checkpoint_id"] - == state_history[0].config["configurable"]["checkpoint_id"] + created_at=AnyStr(), + parent_config=None, + ), + ] + child_state_history = [ + c + for c in app.get_state_history( + {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} ) - # check resuming from interrupt w/ checkpoint_id - interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] - before_interrupt_config = before_interrupt_state_snapshot.config - # going to get to interrupt again here, so the output is None - assert app.invoke(None, before_interrupt_config, debug=True) == { - "my_key": "hi my value" - } - # one more "identical" snapshot than before, at top of list - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + ] + assert child_state_history == [ + StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", } }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + ), + # there should be a single child checkpoint because we only keep + # one child checkpoint per parent checkpoint (in which child ran) + ] + + # check that child snapshot matches id of parent + child_snapshot = child_state_history[0] + assert ( + child_snapshot.config["configurable"]["checkpoint_id"] + == state_history[0].config["configurable"]["checkpoint_id"] + ) + # check resuming from interrupt w/ checkpoint_id + interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] + before_interrupt_config = before_interrupt_state_snapshot.config + # going to get to interrupt again here, so the output is None + assert app.invoke(None, before_interrupt_config, debug=True) == { + "my_key": "hi my value" + } + # one more "identical" snapshot than before, at top of list + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # going to restart from interrupt + interrupt_config = interrupt_state_snapshot.config + assert app.invoke(None, interrupt_config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + assert list(app.get_state_history(config)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # going to restart from interrupt - interrupt_config = interrupt_state_snapshot.config - assert app.invoke(None, interrupt_config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - SqliteSaver.from_conn_string(":memory:"), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - ], - ids=["memory", "sqlite", "postgres"], + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], ) -def test_nested_graph_interrupts_parallel(checkpointer: BaseCheckpointSaver) -> None: - with checkpointer as checkpointer: +def test_nested_graph_interrupts_parallel( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class InnerState(TypedDict): - my_key: Annotated[str, operator.add] - my_other_key: str + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str - def inner_1(state: InnerState): - time.sleep(0.1) - return {"my_key": "got here", "my_other_key": state["my_key"]} + def inner_1(state: InnerState): + time.sleep(0.1) + return {"my_key": "got here", "my_other_key": state["my_key"]} - def inner_2(state: InnerState): - return { - "my_key": " and there", - "my_other_key": state["my_key"], - } + 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") + 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] + class State(TypedDict): + my_key: Annotated[str, operator.add] - def outer_1(state: State): - return {"my_key": " and parallel"} + def outer_1(state: State): + return {"my_key": " and parallel"} - def outer_2(state: State): - return {"my_key": " and back again"} + 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 = 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") + 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) + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": ""}, config, debug=True) == { + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": ""}, config, debug=True) == { + "my_key": "", + } + + assert app.invoke(None, config, debug=True) == { + "my_key": "got here and there and parallel and back again", + } + + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [*app.stream({"my_key": ""}, config)] == [ + # we got to parallel node first + {"outer_1": {"my_key": " and parallel"}}, + ] + assert [*app.stream(None, config)] == [ + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + { "my_key": "", - } - - assert app.invoke(None, config, debug=True) == { + }, + ] + 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", - } + }, + ] - # below combo of assertions is asserting two things - # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) - # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [*app.stream({"my_key": ""}, config)] == [ - # we got to parallel node first - {"outer_1": {"my_key": " and parallel"}}, - ] - assert [*app.stream(None, config)] == [ - {"inner": {"my_key": "got here and there"}}, - {"outer_2": {"my_key": " and back again"}}, - ] + # test 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 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", - }, - ] + # 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", + }, + ] @pytest.mark.skip @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - SqliteSaver.from_conn_string(":memory:"), - PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - ], - ids=["memory", "sqlite", "postgres"], + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], ) -def test_doubly_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None: - with checkpointer as checkpointer: +def test_doubly_nested_graph_interrupts( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(checkpointer_name) - class State(TypedDict): - my_key: str + class State(TypedDict): + my_key: str - class ChildState(TypedDict): - my_key: str + class ChildState(TypedDict): + my_key: str - class GrandChildState(TypedDict): - my_key: str + class GrandChildState(TypedDict): + my_key: str - def grandchild_1(state: ChildState): - return {"my_key": state["my_key"] + " here"} + 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", - } + 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") + 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") + 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_1(state: State): + return {"my_key": "hi " + state["my_key"]} - def parent_2(state: State): - return {"my_key": state["my_key"] + " and back again"} + 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") + 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) + 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) == { + # 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.invoke(None, config, debug=True) == { + }, + ] + 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 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", - }, - ] + }, + ] def test_repeat_condition(snapshot: SnapshotAssertion) -> None: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 020a2c1a8..91ecbc448 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -251,181 +251,159 @@ async def test_step_timeout_on_stream_hang() -> None: @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - NoneContextManager(), - ], - ids=[ - "memory", - "aiosqlite", - "postgres", - "postgres_pipeline", - "none", - ], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) async def test_cancel_graph_astream( - checkpointer: Optional[BaseCheckpointSaver], + request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - async with checkpointer as checkpointer: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): - value: Annotated[int, operator.add] + class State(TypedDict): + value: Annotated[int, operator.add] - class AwhileMaker: - def __init__(self) -> None: - self.reset() + class AwhileMaker: + def __init__(self) -> None: + self.reset() - async def __call__(self, input: State) -> Any: - self.started = True - try: - await asyncio.sleep(1.5) - except asyncio.CancelledError: - self.cancelled = True - raise + async def __call__(self, input: State) -> Any: + self.started = True + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + self.cancelled = True + raise - def reset(self): - self.started = False - self.cancelled = False + def reset(self): + self.started = False + self.cancelled = False - async def alittlewhile(input: State) -> None: - await asyncio.sleep(0.6) - return {"value": 2} + async def alittlewhile(input: State) -> None: + await asyncio.sleep(0.6) + return {"value": 2} - awhile = AwhileMaker() - aparallelwhile = AwhileMaker() - builder = StateGraph(State) - builder.add_node("awhile", awhile) - builder.add_node("aparallelwhile", aparallelwhile) - builder.add_node(alittlewhile) - builder.add_edge(START, "alittlewhile") - builder.add_edge(START, "aparallelwhile") - builder.add_edge("alittlewhile", "awhile") - graph = builder.compile(checkpointer=checkpointer) + awhile = AwhileMaker() + aparallelwhile = AwhileMaker() + builder = StateGraph(State) + builder.add_node("awhile", awhile) + builder.add_node("aparallelwhile", aparallelwhile) + builder.add_node(alittlewhile) + builder.add_edge(START, "alittlewhile") + builder.add_edge(START, "aparallelwhile") + builder.add_edge("alittlewhile", "awhile") + graph = builder.compile(checkpointer=checkpointer) - # test interrupting astream - got_event = False - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - async with aclosing(graph.astream({"value": 1}, thread1)) as stream: - async for chunk in stream: - assert chunk == {"alittlewhile": {"value": 2}} - got_event = True - break + # test interrupting astream + got_event = False + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + async with aclosing(graph.astream({"value": 1}, thread1)) as stream: + async for chunk in stream: + assert chunk == {"alittlewhile": {"value": 2}} + got_event = True + break - assert got_event + assert got_event - # node aparallelwhile should start, but be cancelled - assert aparallelwhile.started is True - assert aparallelwhile.cancelled is True + # node aparallelwhile should start, but be cancelled + assert aparallelwhile.started is True + assert aparallelwhile.cancelled is True - # node "awhile" should never start - assert awhile.started is False + # node "awhile" should never start + assert awhile.started is False - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ( - "aparallelwhile", - "alittlewhile", - ) - assert state.metadata == {"source": "loop", "step": 0, "writes": None} + # checkpoint with output of "alittlewhile" should not be saved + if checkpointer is not None: + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ( + "aparallelwhile", + "alittlewhile", + ) + assert state.metadata == {"source": "loop", "step": 0, "writes": None} @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - NoneContextManager(), - ], - ids=[ - "memory", - "aiosqlite", - "postgres", - "postgres_pipeline", - "none", - ], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe", None], ) async def test_cancel_graph_astream_events_v2( - checkpointer: Optional[BaseCheckpointSaver], + request: pytest.FixtureRequest, checkpointer_name: Optional[str] ) -> None: - async with checkpointer as checkpointer: + checkpointer = ( + request.getfixturevalue(f"checkpointer_{checkpointer_name}") + if checkpointer_name + else None + ) - class State(TypedDict): - value: int + class State(TypedDict): + value: int - class AwhileMaker: - def __init__(self) -> None: - self.reset() + class AwhileMaker: + def __init__(self) -> None: + self.reset() - async def __call__(self, input: State) -> Any: - self.started = True - try: - await asyncio.sleep(1.5) - except asyncio.CancelledError: - self.cancelled = True - raise + async def __call__(self, input: State) -> Any: + self.started = True + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + self.cancelled = True + raise - def reset(self): - self.started = False - self.cancelled = False + def reset(self): + self.started = False + self.cancelled = False - async def alittlewhile(input: State) -> None: - await asyncio.sleep(0.6) - return {"value": 2} + async def alittlewhile(input: State) -> None: + await asyncio.sleep(0.6) + return {"value": 2} - awhile = AwhileMaker() - anotherwhile = AwhileMaker() - builder = StateGraph(State) - builder.add_node(alittlewhile) - builder.add_node("awhile", awhile) - builder.add_node("anotherwhile", anotherwhile) - builder.add_edge(START, "alittlewhile") - builder.add_edge("alittlewhile", "awhile") - builder.add_edge("awhile", "anotherwhile") - graph = builder.compile(checkpointer=checkpointer) + awhile = AwhileMaker() + anotherwhile = AwhileMaker() + builder = StateGraph(State) + builder.add_node(alittlewhile) + builder.add_node("awhile", awhile) + builder.add_node("anotherwhile", anotherwhile) + builder.add_edge(START, "alittlewhile") + builder.add_edge("alittlewhile", "awhile") + builder.add_edge("awhile", "anotherwhile") + graph = builder.compile(checkpointer=checkpointer) - # test interrupting astream_events v2 - got_event = False - thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} - async with aclosing( - graph.astream_events({"value": 1}, thread2, version="v2") - ) as stream: - async for chunk in stream: - if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: - got_event = True - assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} - break + # test interrupting astream_events v2 + got_event = False + thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} + async with aclosing( + graph.astream_events({"value": 1}, thread2, version="v2") + ) as stream: + async for chunk in stream: + if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: + got_event = True + assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} + break - # did break - assert got_event + # did break + assert got_event - # node "awhile" maybe starts (impl detail of astream_events) - # if it does start, it must be cancelled - if awhile.started: - assert awhile.cancelled is True + # node "awhile" maybe starts (impl detail of astream_events) + # if it does start, it must be cancelled + if awhile.started: + assert awhile.cancelled is True - # node "anotherwhile" should never start - assert anotherwhile.started is False + # node "anotherwhile" should never start + assert anotherwhile.started is False - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread2) - assert state is not None - assert state.values == {"value": 2} - assert state.next == ("awhile",) - assert state.metadata == { - "source": "loop", - "step": 1, - "writes": {"alittlewhile": {"value": 2}}, - } + # checkpoint with output of "alittlewhile" should not be saved + if checkpointer is not None: + state = await graph.aget_state(thread2) + assert state is not None + assert state.values == {"value": 2} + assert state.next == ("awhile",) + assert state.metadata == { + "source": "loop", + "step": 1, + "writes": {"alittlewhile": {"value": 2}}, + } async def test_node_schemas_custom_output() -> None: @@ -701,396 +679,379 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) async def test_invoke_two_processes_in_out_interrupt( - checkpointer: BaseCheckpointSaver, mocker: MockerFixture + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: - async with checkpointer as checkpointer: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert await app.ainvoke(2, thread1) is None + + # inbox == 3 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert await app.ainvoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1) is None + + # inbox == 21 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert await app.ainvoke(3, thread1) is None + assert await app.ainvoke(None, thread1) == 5 + + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2) is None + + # inbox == 21 + snapshot = await app.aget_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + await app.aupdate_state(thread2, 25, as_node="one") + assert await app.ainvoke(None, thread2) == 26 + + # no pending tasks + snapshot = await app.aget_state(thread2) + assert snapshot.next == () + + # list history + history = [c async for c in app.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - interrupt_after_nodes=["one"], - ) - thread1 = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} + metadata={"source": "loop", "step": 6, "writes": 5}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": None}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[7].config, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, + created_at=AnyStr(), + parent_config=None, + ), + ] - # start execution, stop at inbox - assert await app.ainvoke(2, thread1) is None + # forking from any previous checkpoint w/out forking should do nothing + assert [ + c async for c in app.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in app.astream(None, history[1].config, stream_mode="updates") + ] == [] + assert [ + c async for c in app.astream(None, history[2].config, stream_mode="updates") + ] == [] - # inbox == 3 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = await app.aupdate_state(history[0].config, None) + assert [ + c async for c in app.astream(None, fork_config, stream_mode="updates") + ] == [] - # resume execution, finish - assert await app.ainvoke(None, thread1) == 4 + fork_config = await app.aupdate_state(history[1].config, None) + assert [c async for c in app.astream(None, fork_config, stream_mode="updates")] == [ + {"two": {"output": 5}} + ] - # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None - - # inbox == 21 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 - - # send a new value in, interrupting the previous execution - assert await app.ainvoke(3, thread1) is None - assert await app.ainvoke(None, thread1) == 5 - - # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) is None - - # inbox == 21 - snapshot = await app.aget_state(thread2) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) - - # update the state, resume - await app.aupdate_state(thread2, 25, as_node="one") - assert await app.ainvoke(None, thread2) == 26 - - # no pending tasks - snapshot = await app.aget_state(thread2) - assert snapshot.next == () - - # list history - history = [c async for c in app.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 6, "writes": 5}, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 5, "writes": None}, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": 4, "writes": 3}, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 3, "writes": None}, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": 2, "writes": 20}, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 1, "writes": 4}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=history[7].config, - ), - StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": -1, "writes": 2}, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # forking from any previous checkpoint w/out forking should do nothing - assert [ - c async for c in app.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in app.astream(None, history[1].config, stream_mode="updates") - ] == [] - assert [ - c async for c in app.astream(None, history[2].config, stream_mode="updates") - ] == [] - - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = await app.aupdate_state(history[0].config, None) - assert [ - c async for c in app.astream(None, fork_config, stream_mode="updates") - ] == [] - - fork_config = await app.aupdate_state(history[1].config, None) - assert [ - c async for c in app.astream(None, fork_config, stream_mode="updates") - ] == [{"two": {"output": 5}}] - - fork_config = await app.aupdate_state(history[2].config, None) - assert [ - c async for c in app.astream(None, fork_config, stream_mode="updates") - ] == [{"one": {"inbox": 4}}] + fork_config = await app.aupdate_state(history[2].config, None) + assert [c async for c in app.astream(None, fork_config, stream_mode="updates")] == [ + {"one": {"inbox": 4}} + ] @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) async def test_fork_always_re_runs_nodes( - checkpointer: BaseCheckpointSaver, mocker: MockerFixture + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: - async with checkpointer as checkpointer: - add_one = mocker.Mock(side_effect=lambda _: 1) + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + add_one = mocker.Mock(side_effect=lambda _: 1) - builder = StateGraph(Annotated[int, operator.add]) - builder.add_node("add_one", add_one) - builder.add_edge(START, "add_one") - builder.add_conditional_edges( - "add_one", lambda cnt: "add_one" if cnt < 6 else END - ) - graph = builder.compile(checkpointer=checkpointer) + builder = StateGraph(Annotated[int, operator.add]) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END) + graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} + thread1 = {"configurable": {"thread_id": "1"}} - # start execution, stop at inbox - assert [ - c - async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) - ] == [ - ("values", 1), - ("updates", {"add_one": 1}), - ("values", 2), - ("updates", {"add_one": 1}), - ("values", 3), - ("updates", {"add_one": 1}), - ("values", 4), - ("updates", {"add_one": 1}), - ("values", 5), - ("updates", {"add_one": 1}), - ("values", 6), - ] + # start execution, stop at inbox + assert [ + c async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) + ] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] - # list history - history = [c async for c in graph.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values=6, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values=5, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values=4, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values=3, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values=2, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values=1, - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values=0, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "input", "step": -1, "writes": 1}, - created_at=AnyStr(), - parent_config=None, - ), - ] + # list history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values=5, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values=4, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values=3, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values=2, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values=1, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values=0, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 1}, + created_at=AnyStr(), + parent_config=None, + ), + ] - # forking from any previous checkpoint w/out forking should do nothing - assert [ - c - async for c in graph.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c - async for c in graph.astream(None, history[1].config, stream_mode="updates") - ] == [] + # forking from any previous checkpoint w/out forking should do nothing + assert [ + c async for c in graph.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in graph.astream(None, history[1].config, stream_mode="updates") + ] == [] - # forking and re-running from any prev checkpoint should re-run nodes - fork_config = await graph.aupdate_state(history[0].config, None) - assert [ - c async for c in graph.astream(None, fork_config, stream_mode="updates") - ] == [] + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = await graph.aupdate_state(history[0].config, None) + assert [ + c async for c in graph.astream(None, fork_config, stream_mode="updates") + ] == [] - fork_config = await graph.aupdate_state(history[1].config, None) - assert [ - c async for c in graph.astream(None, fork_config, stream_mode="updates") - ] == [{"add_one": 1}] + fork_config = await graph.aupdate_state(history[1].config, None) + assert [ + c async for c in graph.astream(None, fork_config, stream_mode="updates") + ] == [{"add_one": 1}] - fork_config = await graph.aupdate_state(history[2].config, None) - assert [ - c async for c in graph.astream(None, fork_config, stream_mode="updates") - ] == [ - {"add_one": 1}, - {"add_one": 1}, - ] + fork_config = await graph.aupdate_state(history[2].config, None) + assert [ + c async for c in graph.astream(None, fork_config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: @@ -1422,88 +1383,84 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) -async def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None: - async with checkpointer as checkpointer: +async def test_pending_writes_resume( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): - value: Annotated[int, operator.add] + class State(TypedDict): + value: Annotated[int, operator.add] - class AwhileMaker: - def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None: - self.sleep = sleep - self.rtn = rtn - self.reset() + class AwhileMaker: + def __init__(self, sleep: float, rtn: Union[Dict, Exception]) -> None: + self.sleep = sleep + self.rtn = rtn + self.reset() - async def __call__(self, input: State) -> Any: - self.calls += 1 - await asyncio.sleep(self.sleep) - if isinstance(self.rtn, Exception): - raise self.rtn - else: - return self.rtn + async def __call__(self, input: State) -> Any: + self.calls += 1 + await asyncio.sleep(self.sleep) + if isinstance(self.rtn, Exception): + raise self.rtn + else: + return self.rtn - def reset(self): - self.calls = 0 + def reset(self): + self.calls = 0 - one = AwhileMaker(0.2, {"value": 2}) - two = AwhileMaker(0.6, ValueError("I'm not good")) - builder = StateGraph(State) - builder.add_node("one", one) - builder.add_node("two", two) - builder.add_edge(START, "one") - builder.add_edge(START, "two") - graph = builder.compile(checkpointer=checkpointer) + one = AwhileMaker(0.2, {"value": 2}) + two = AwhileMaker(0.6, ValueError("I'm not good")) + builder = StateGraph(State) + builder.add_node("one", one) + builder.add_node("two", two) + builder.add_edge(START, "one") + builder.add_edge(START, "two") + graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - with pytest.raises(ValueError, match="I'm not good"): - await graph.ainvoke({"value": 1}, thread1) + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ValueError, match="I'm not good"): + await graph.ainvoke({"value": 1}, thread1) - # both nodes should have been called once - assert one.calls == 1 - assert two.calls == 1 + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 1 - # latest checkpoint should be before nodes "one", "two" - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") - assert state.metadata == {"source": "loop", "step": 0, "writes": None} - # should contain pending write of "one" - checkpoint = await checkpointer.aget_tuple(thread1) - assert checkpoint is not None - assert checkpoint.pending_writes == [ - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - ] - # both pending writes come from same task - assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0] + # latest checkpoint should be before nodes "one", "two" + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + assert state.metadata == {"source": "loop", "step": 0, "writes": None} + # should contain pending write of "one" + checkpoint = await checkpointer.aget_tuple(thread1) + assert checkpoint is not None + assert checkpoint.pending_writes == [ + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + ] + # both pending writes come from same task + assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0] - # resume execution - with pytest.raises(ValueError, match="I'm not good"): - await graph.ainvoke(None, thread1) + # resume execution + with pytest.raises(ValueError, match="I'm not good"): + await graph.ainvoke(None, thread1) - # node "one" succeeded previously, so shouldn't be called again - assert one.calls == 1 - # node "two" should have been called once again - assert two.calls == 2 + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 2 - # confirm no new checkpoints saved - state_two = await graph.aget_state(thread1) - assert state_two == state + # confirm no new checkpoints saved + state_two = await graph.aget_state(thread1) + assert state_two == state - # resume execution, without exception - two.rtn = {"value": 3} - # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke(None, thread1) == {"value": 6} + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert await graph.ainvoke(None, thread1) == {"value": 6} async def test_cond_edge_after_send() -> None: @@ -6147,1450 +6104,1414 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.repeat(10) @pytest.mark.parametrize( - "checkpointer_fct", - [ - lambda: MemorySaverAssertImmutable(put_sleep=0.2), - lambda: AsyncSqliteSaver.from_conn_string(":memory:"), - lambda: AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - lambda: AsyncPostgresSaver.from_conn_string( - DEFAULT_POSTGRES_URI, pipeline=True - ), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) async def test_nested_graph_interrupts( - checkpointer_fct: Callable[[], BaseCheckpointSaver], + request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - async with checkpointer_fct() as checkpointer: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class InnerState(TypedDict): - my_key: str - my_other_key: str + 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_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"], - } + 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") + 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 + class State(TypedDict): + my_key: str - async def outer_1(state: State): - return {"my_key": "hi " + state["my_key"]} + 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"} + 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") + 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) + 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) == { + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + {"outer_1": {"my_key": "hi my value"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"inner": {"my_key": "hi my value here and there"}}, + {"outer_2": {"my_key": "hi my value here and there and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") + ] == [ + { + "my_key": "my value", + }, + { "my_key": "hi my value", - } - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": 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) == { + }, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { "my_key": "hi my value here and there and back again", - } - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ - {"outer_1": {"my_key": "hi my value"}}, - ] - assert [c async for c in app.astream(None, config)] == [ - {"inner": {"my_key": "hi my value here and there"}}, - {"outer_2": {"my_key": "hi my value here and there and back again"}}, - ] + }, + ] - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream( - {"my_key": "my value"}, config, stream_mode="values" - ) - ] == [ - { - "my_key": "my value", + # test interrupts BEFORE the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c + async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") + ] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value", + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, }, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value here and there and back again", + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - ] + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [c async for c in app.astream(None, config, stream_mode="values")] == [] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + { + "my_key": "hi my value here and there and back again", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # 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", + # test interrupts AFTER the node w/ interrupts + app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c + async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") + ] == [ + { + "my_key": "my value", + }, + { + "my_key": "hi my value", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value", + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, }, - ] - 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": 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", + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - { - "my_key": "hi my value here and there and back again", + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + { + "my_key": "hi my value here and there and back again", + }, + ] + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": 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": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # test interrupts AFTER the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c - async for c in app.astream( - {"my_key": "my value"}, config, stream_mode="values" - ) - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there and back again", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + # test restarting from checkpoint_id + config = {"configurable": {"thread_id": "6"}} + app = graph.compile(checkpointer=checkpointer) + await app.ainvoke({"my_key": "my value"}, config, debug=True) - # test restarting from checkpoint_id - config = {"configurable": {"thread_id": "6"}} - app = graph.compile(checkpointer=checkpointer) - await app.ainvoke({"my_key": "my value"}, config, debug=True) + state_history = [c async for c in app.aget_state_history(config)] + assert state_history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - state_history = [c async for c in app.aget_state_history(config)] - assert state_history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - child_state_history = [ - c - async for c in app.aget_state_history( - {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} - ) - ] - assert child_state_history == [ - StateSnapshot( - values={"my_key": "hi my value here"}, - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": 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", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - # there should be a single child checkpoint because we only keep - # one child checkpoint per parent checkpoint (in which child ran) - ] - - # check that child snapshot matches id of parent - child_snapshot = child_state_history[0] - assert ( - child_snapshot.config["configurable"]["checkpoint_id"] - == state_history[0].config["configurable"]["checkpoint_id"] + child_state_history = [ + c + async for c in app.aget_state_history( + {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} ) - # check resuming from interrupt w/ checkpoint_id - interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] - before_interrupt_config = before_interrupt_state_snapshot.config - # going to get to interrupt again here - assert await app.ainvoke(None, before_interrupt_config, debug=True) == { - "my_key": "hi my value" - } - # one more "identical" snapshot than before, at top of list - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + ] + assert child_state_history == [ + StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", } }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + ), + # there should be a single child checkpoint because we only keep + # one child checkpoint per parent checkpoint (in which child ran) + ] + + # check that child snapshot matches id of parent + child_snapshot = child_state_history[0] + assert ( + child_snapshot.config["configurable"]["checkpoint_id"] + == state_history[0].config["configurable"]["checkpoint_id"] + ) + # check resuming from interrupt w/ checkpoint_id + interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] + before_interrupt_config = before_interrupt_state_snapshot.config + # going to get to interrupt again here + assert await app.ainvoke(None, before_interrupt_config, debug=True) == { + "my_key": "hi my value" + } + # one more "identical" snapshot than before, at top of list + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # going to resume from interrupt + interrupt_config = interrupt_state_snapshot.config + assert (await app.ainvoke(None, interrupt_config, debug=True)) == { + "my_key": "hi my value here and there and back again", + } + assert [s async for s in app.aget_state_history(config)] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # going to resume from interrupt - interrupt_config = interrupt_state_snapshot.config - assert (await app.ainvoke(None, interrupt_config, debug=True)) == { - "my_key": "hi my value here and there and back again", - } - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("outer_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": 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", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": 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", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) async def test_nested_graph_interrupts_parallel( - checkpointer: BaseCheckpointSaver, + request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - async with checkpointer as checkpointer: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class InnerState(TypedDict): - my_key: Annotated[str, operator.add] - my_other_key: str + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str - async def inner_1(state: InnerState): - await asyncio.sleep(0.1) - return {"my_key": "got here", "my_other_key": state["my_key"]} + async def inner_1(state: InnerState): + await asyncio.sleep(0.1) + return {"my_key": "got here", "my_other_key": state["my_key"]} - async def inner_2(state: InnerState): - return { - "my_key": " and there", - "my_other_key": state["my_key"], - } + 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") + 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] + class State(TypedDict): + my_key: Annotated[str, operator.add] - async def outer_1(state: State): - return {"my_key": " and parallel"} + async def outer_1(state: State): + return {"my_key": " and parallel"} - async def outer_2(state: State): - return {"my_key": " and back again"} + 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 = 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") + 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) + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + "my_key": "", + } + + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "got here and there and parallel and back again", + } + + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": ""}, config)] == [ + # we got to parallel node first + {"outer_1": {"my_key": " and parallel"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + { "my_key": "", - } - - assert await app.ainvoke(None, config, debug=True) == { + }, + ] + 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", - } + }, + ] - # below combo of assertions is asserting two things - # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) - # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": ""}, config)] == [ - # we got to parallel node first - {"outer_1": {"my_key": " and parallel"}}, - ] - assert [c async for c in app.astream(None, config)] == [ - {"inner": {"my_key": "got here and there"}}, - {"outer_2": {"my_key": " and back again"}}, - ] + # # test 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 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", - }, - ] + # 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", + }, + ] @pytest.mark.skip @pytest.mark.parametrize( - "checkpointer", - [ - MemorySaverAssertImmutable(), - AsyncSqliteSaver.from_conn_string(":memory:"), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI), - AsyncPostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI, pipeline=True), - ], - ids=["memory", "sqlite", "postgres", "postgres_pipeline"], + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], ) async def test_doubly_nested_graph_interrupts( - checkpointer: BaseCheckpointSaver, + request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - async with checkpointer as checkpointer: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): - my_key: str + class State(TypedDict): + my_key: str - class ChildState(TypedDict): - my_key: str + class ChildState(TypedDict): + my_key: str - class GrandChildState(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_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} - async def grandchild_2(state: ChildState): - return { - "my_key": state["my_key"] + " and there", - } + 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") + 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") + 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_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"} + 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") + 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) + 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) == { + # 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 await app.ainvoke(None, config, debug=True) == { + }, + ] + 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 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", - }, - ] + }, + ] async def test_checkpoint_metadata() -> None: From 6028eb02c9d4e9f5c1e397b1112c0bb3d358d19b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 7 Aug 2024 11:49:46 -0700 Subject: [PATCH 2/3] Lint --- libs/langgraph/tests/test_pregel_async.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 91ecbc448..d7303eb28 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9,7 +9,6 @@ from typing import ( Any, AsyncGenerator, AsyncIterator, - Callable, Dict, Generator, List, @@ -41,14 +40,12 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( - BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.constants import Send from langgraph.errors import InvalidUpdateError @@ -63,7 +60,6 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr -from tests.conftest import DEFAULT_POSTGRES_URI from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, From cb21a3547ebffb191be21e23548ddbc989beb32e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 7 Aug 2024 12:04:11 -0700 Subject: [PATCH 3/3] Skip async postgres tests in 3.9 --- libs/langgraph/tests/conftest.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index f0c85f9a6..1219ff0f1 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,4 +1,5 @@ import asyncio +import sys from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager, contextmanager from typing import AsyncIterator, Iterator, TypeVar @@ -123,6 +124,8 @@ def checkpointer_postgres_pipe(): @pytest.fixture(scope="function") def checkpointer_postgres_aio(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") with agen_to_gen(_checkpointer_postgres_aio()) as checkpointer: yield checkpointer @@ -152,6 +155,8 @@ async def _checkpointer_postgres_aio(): @pytest.fixture(scope="function") def checkpointer_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") with agen_to_gen(_checkpointer_postgres_aio_pipe()) as checkpointer: yield checkpointer