From af961e279baef505075eff3bae5ebd105a83c94a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 09:59:43 -0700 Subject: [PATCH 1/2] More idiomatic sync/async checkpointer fixtures in pytest --- libs/langgraph/tests/conftest.py | 328 ++-- libs/langgraph/tests/conftest_checkpointer.py | 245 +++ .../tests/test_checkpoint_migration.py | 205 ++- libs/langgraph/tests/test_interruption.py | 48 +- libs/langgraph/tests/test_large_cases.py | 22 +- .../langgraph/tests/test_large_cases_async.py | 380 +++-- libs/langgraph/tests/test_pregel.py | 63 +- libs/langgraph/tests/test_pregel_async.py | 1423 ++++++++--------- 8 files changed, 1366 insertions(+), 1348 deletions(-) create mode 100644 libs/langgraph/tests/conftest_checkpointer.py diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 0796112e3..504cb5785 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,29 +1,32 @@ -import sys from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager from typing import Optional -from uuid import UUID, uuid4 +from uuid import UUID import pytest from langchain_core import __version__ as core_version from packaging import version -from psycopg import AsyncConnection, Connection -from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture from langgraph.cache.base import BaseCache from langgraph.cache.memory import InMemoryCache from langgraph.cache.sqlite import SqliteCache from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver -from langgraph.checkpoint.postgres.aio import ( - AsyncPostgresSaver, - AsyncShallowPostgresSaver, -) -from langgraph.checkpoint.serde.encrypted import EncryptedSerializer -from langgraph.checkpoint.sqlite import SqliteSaver -from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.store.base import BaseStore +from tests.conftest_checkpointer import ( + _checkpointer_memory, + _checkpointer_postgres, + _checkpointer_postgres_aio, + _checkpointer_postgres_aio_pipe, + _checkpointer_postgres_aio_pool, + _checkpointer_postgres_aio_shallow, + _checkpointer_postgres_pipe, + _checkpointer_postgres_pool, + _checkpointer_postgres_shallow, + _checkpointer_sqlite, + _checkpointer_sqlite_aes, + _checkpointer_sqlite_aio, +) from tests.conftest_store import ( _store_memory, _store_postgres, @@ -57,219 +60,54 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: return mocker.patch("uuid.uuid4", side_effect=side_effect) -# checkpointer fixtures +@pytest.fixture(params=[True, False]) +def checkpoint_during(request: pytest.FixtureRequest) -> bool: + return request.param + + +# --- start of deprecated fixtures --- @pytest.fixture(scope="function") def checkpointer_memory(): - from tests.memory_assert import MemorySaverAssertImmutable - - yield MemorySaverAssertImmutable() + with _checkpointer_memory() as checkpointer: + yield checkpointer @pytest.fixture(scope="function") def checkpointer_sqlite(): - with SqliteSaver.from_conn_string(":memory:") as checkpointer: + with _checkpointer_sqlite() as checkpointer: yield checkpointer @pytest.fixture(scope="function") def checkpointer_sqlite_aes(): - with SqliteSaver.from_conn_string(":memory:") as checkpointer: - checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes( - key=b"1234567890123456" - ) - yield checkpointer - - -@asynccontextmanager -async def _checkpointer_sqlite_aio(): - async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: + with _checkpointer_sqlite_aes() 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: - # 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}") + with _checkpointer_postgres() as checkpointer: + yield checkpointer @pytest.fixture(scope="function") def checkpointer_postgres_shallow(): - 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 ShallowPostgresSaver.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}") + with _checkpointer_postgres_shallow() as checkpointer: + yield checkpointer @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}") + with _checkpointer_postgres_pipe() as checkpointer: + yield checkpointer @pytest.fixture(scope="function") def checkpointer_postgres_pool(): - 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 ConnectionPool( - DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} - ) as pool: - checkpointer = PostgresSaver(pool) - checkpointer.setup() - yield checkpointer - finally: - # drop unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"DROP DATABASE {database}") - - -@asynccontextmanager -async def _checkpointer_postgres_aio(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - 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}") - - -@asynccontextmanager -async def _checkpointer_postgres_aio_shallow(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - 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 AsyncShallowPostgresSaver.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}") - - -@asynccontextmanager -async def _checkpointer_postgres_aio_pipe(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - 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}") - - -@asynccontextmanager -async def _checkpointer_postgres_aio_pool(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - 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 AsyncConnectionPool( - DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} - ) as pool: - checkpointer = AsyncPostgresSaver(pool) - 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}") + with _checkpointer_postgres_pool() as checkpointer: + yield checkpointer @asynccontextmanager @@ -279,9 +117,8 @@ async def awith_checkpointer( if checkpointer_name is None: yield None elif checkpointer_name == "memory": - from tests.memory_assert import MemorySaverAssertImmutable - - yield MemorySaverAssertImmutable() + with _checkpointer_memory() as checkpointer: + yield checkpointer elif checkpointer_name == "sqlite_aio": async with _checkpointer_sqlite_aio() as checkpointer: yield checkpointer @@ -301,6 +138,9 @@ async def awith_checkpointer( raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") +# --- end of deprecated fixtures --- + + @pytest.fixture(scope="function", params=["sqlite", "memory"]) def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: if request.param == "sqlite": @@ -359,32 +199,90 @@ async def async_store(request: pytest.FixtureRequest) -> AsyncIterator[BaseStore raise NotImplementedError(f"Unknown store {store_name}") -SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"] -REGULAR_CHECKPOINTERS_SYNC = [ +@pytest.fixture( + scope="function", + params=[ + "memory", + "sqlite", + "sqlite_aes", + "postgres", + "postgres_pipe", + "postgres_pool", + ], +) +def sync_checkpointer( + request: pytest.FixtureRequest, +) -> Iterator[BaseCheckpointSaver]: + checkpointer_name = request.param + if checkpointer_name == "memory": + with _checkpointer_memory() as checkpointer: + yield checkpointer + elif checkpointer_name == "sqlite": + with _checkpointer_sqlite() as checkpointer: + yield checkpointer + elif checkpointer_name == "sqlite_aes": + with _checkpointer_sqlite_aes() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres": + with _checkpointer_postgres() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_pipe": + with _checkpointer_postgres_pipe() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_pool": + with _checkpointer_postgres_pool() as checkpointer: + yield checkpointer + else: + raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") + + +@pytest.fixture( + scope="function", + params=[ + "memory", + "sqlite_aio", + "postgres_aio", + "postgres_aio_pipe", + "postgres_aio_pool", + ], +) +async def async_checkpointer( + request: pytest.FixtureRequest, +) -> AsyncIterator[BaseCheckpointSaver]: + checkpointer_name = request.param + if checkpointer_name == "memory": + with _checkpointer_memory() as checkpointer: + yield checkpointer + elif checkpointer_name == "sqlite_aio": + async with _checkpointer_sqlite_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio": + async with _checkpointer_postgres_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pipe": + async with _checkpointer_postgres_aio_pipe() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pool": + async with _checkpointer_postgres_aio_pool() as checkpointer: + yield checkpointer + else: + raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") + + +ALL_CHECKPOINTERS_SYNC = [ "memory", "sqlite", + "sqlite_aes", "postgres", "postgres_pipe", "postgres_pool", - "sqlite_aes", + "postgres_shallow", ] -ALL_CHECKPOINTERS_SYNC = [ - *REGULAR_CHECKPOINTERS_SYNC, - *SHALLOW_CHECKPOINTERS_SYNC, -] -SHALLOW_CHECKPOINTERS_ASYNC = ["postgres_aio_shallow"] -REGULAR_CHECKPOINTERS_ASYNC = [ +ALL_CHECKPOINTERS_ASYNC = [ "memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe", "postgres_aio_pool", -] -ALL_CHECKPOINTERS_ASYNC = [ - *REGULAR_CHECKPOINTERS_ASYNC, - *SHALLOW_CHECKPOINTERS_ASYNC, -] -ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ - *ALL_CHECKPOINTERS_ASYNC, - None, + "postgres_aio_shallow", ] diff --git a/libs/langgraph/tests/conftest_checkpointer.py b/libs/langgraph/tests/conftest_checkpointer.py new file mode 100644 index 000000000..e43773586 --- /dev/null +++ b/libs/langgraph/tests/conftest_checkpointer.py @@ -0,0 +1,245 @@ +import sys +from contextlib import asynccontextmanager, contextmanager +from uuid import uuid4 + +import pytest +from psycopg import AsyncConnection, Connection +from psycopg_pool import AsyncConnectionPool, ConnectionPool + +from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver +from langgraph.checkpoint.postgres.aio import ( + AsyncPostgresSaver, + AsyncShallowPostgresSaver, +) +from langgraph.checkpoint.serde.encrypted import EncryptedSerializer +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/" + + +@contextmanager +def _checkpointer_memory(): + yield MemorySaverAssertImmutable() + + +@contextmanager +def _checkpointer_sqlite(): + with SqliteSaver.from_conn_string(":memory:") as checkpointer: + yield checkpointer + + +@contextmanager +def _checkpointer_sqlite_aes(): + with SqliteSaver.from_conn_string(":memory:") as checkpointer: + checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes( + key=b"1234567890123456" + ) + yield checkpointer + + +@contextmanager +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: + # 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}") + + +@contextmanager +def _checkpointer_postgres_shallow(): + 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 ShallowPostgresSaver.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}") + + +@contextmanager +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}") + + +@contextmanager +def _checkpointer_postgres_pool(): + 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 ConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = PostgresSaver(pool) + checkpointer.setup() + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _checkpointer_sqlite_aio(): + async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_postgres_aio(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + 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}") + + +@asynccontextmanager +async def _checkpointer_postgres_aio_shallow(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + 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 AsyncShallowPostgresSaver.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}") + + +@asynccontextmanager +async def _checkpointer_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + 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}") + + +@asynccontextmanager +async def _checkpointer_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + 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 AsyncConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = AsyncPostgresSaver(pool) + 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}") + + +__all__ = [ + "_checkpointer_memory", + "_checkpointer_sqlite", + "_checkpointer_sqlite_aes", + "_checkpointer_postgres", + "_checkpointer_postgres_shallow", + "_checkpointer_postgres_pipe", + "_checkpointer_postgres_pool", + "_checkpointer_sqlite_aio", + "_checkpointer_postgres_aio", + "_checkpointer_postgres_aio_shallow", + "_checkpointer_postgres_aio_pipe", + "_checkpointer_postgres_aio_pool", +] diff --git a/libs/langgraph/tests/test_checkpoint_migration.py b/libs/langgraph/tests/test_checkpoint_migration.py index 1ac86adba..caba88060 100644 --- a/libs/langgraph/tests/test_checkpoint_migration.py +++ b/libs/langgraph/tests/test_checkpoint_migration.py @@ -17,11 +17,6 @@ from langgraph.types import Command, Interrupt, PregelTask, StateSnapshot, inter from langgraph.utils.config import patch_configurable from tests.any_int import AnyInt from tests.any_str import AnyDict, AnyObject, AnyStr -from tests.conftest import ( - REGULAR_CHECKPOINTERS_ASYNC, - REGULAR_CHECKPOINTERS_SYNC, - awith_checkpointer, -) pytestmark = pytest.mark.anyio @@ -1593,16 +1588,11 @@ def test_migrate_checkpoints(source: str, target: str) -> None: @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_latest_checkpoint_state_graph( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - builder = make_state_graph() - app = builder.compile(checkpointer=checkpointer) + app = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} assert [*app.stream({"query": "what is weather in sf"}, config)] == [ @@ -1638,50 +1628,49 @@ def test_latest_checkpoint_state_graph( @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_latest_checkpoint_state_graph_async(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - builder = make_state_graph() - app = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} +async def test_latest_checkpoint_state_graph_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + builder = make_state_graph() + app = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} - assert [ - c async for c in app.astream({"query": "what is weather in sf"}, config) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - { - "__interrupt__": ( - Interrupt( - value="", - resumable=True, - ns=[AnyStr("qa:")], - ), - ) - }, - ] + assert [ + c async for c in app.astream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + { + "__interrupt__": ( + Interrupt( + value="", + resumable=True, + ns=[AnyStr("qa:")], + ), + ) + }, + ] - assert [c async for c in app.astream(Command(resume=""), config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app.astream(Command(resume=""), config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - # check history with current checkpoints matches expected history - history = [c async for c in app.aget_state_history(config)] - expected_history = get_expected_history() - assert len(history) == len(expected_history) - assert history[0] == expected_history[0] - assert history[1] == expected_history[1] - assert history[2] == expected_history[2] - assert history[3] == expected_history[3] - assert history[4] == expected_history[4] - assert history[5] == expected_history[5] + # check history with current checkpoints matches expected history + history = [c async for c in app.aget_state_history(config)] + expected_history = get_expected_history() + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_saved_checkpoint_state_graph( request: pytest.FixtureRequest, checkpointer_name: str, @@ -1753,71 +1742,65 @@ def test_saved_checkpoint_state_graph( @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_saved_checkpoint_state_graph_async( - checkpointer_name: str, + async_checkpointer: BaseCheckpointSaver, checkpoint_version: str, ) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - builder = make_state_graph() - app = builder.compile(checkpointer=checkpointer) + builder = make_state_graph() + app = builder.compile(checkpointer=async_checkpointer) - thread1 = "1" - config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} + thread1 = "1" + config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} - # save checkpoints - parent_id: Optional[str] = None - for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]): - grouped_writes = defaultdict(list) - for write in checkpoint.pending_writes: - grouped_writes[write[0]].append(write[1:]) - for tid, group in grouped_writes.items(): - await checkpointer.aput_writes(checkpoint.config, group, tid) - await checkpointer.aput( - patch_configurable(config, {"checkpoint_id": parent_id}), - checkpoint.checkpoint, - checkpoint.metadata, - checkpoint.checkpoint["channel_versions"], - ) - parent_id = checkpoint.checkpoint["id"] - - # load history - history = [c async for c in app.aget_state_history(config)] - # check history with saved checkpoints matches expected history - exc_task_results: int = 0 - if checkpoint_version == "2-start:*": - exc_task_results = 1 - elif checkpoint_version == "2-quadratic": - exc_task_results = 2 - expected_history = get_expected_history(exc_task_results=exc_task_results) - assert len(history) == len(expected_history) - assert history[0] == expected_history[0] - assert history[1] == expected_history[1] - assert history[2] == expected_history[2] - assert history[3] == expected_history[3] - assert history[4] == expected_history[4] - assert history[5] == expected_history[5] - - # resume from 2nd to latest checkpoint - assert [ - c async for c in app.astream(Command(resume=""), history[1].config) - ] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - # new checkpoint should match the latest checkpoint in history - latest_state = await app.aget_state(config) - assert ( - StateSnapshot( - values=latest_state.values, - next=latest_state.next, - config=patch_configurable( - latest_state.config, {"checkpoint_id": AnyStr()} - ), - metadata=AnyDict(latest_state.metadata), - created_at=AnyStr(), - parent_config=latest_state.parent_config, - tasks=latest_state.tasks, - interrupts=latest_state.interrupts, - ) - == history[0] + # save checkpoints + parent_id: Optional[str] = None + for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]): + grouped_writes = defaultdict(list) + for write in checkpoint.pending_writes: + grouped_writes[write[0]].append(write[1:]) + for tid, group in grouped_writes.items(): + await async_checkpointer.aput_writes(checkpoint.config, group, tid) + await async_checkpointer.aput( + patch_configurable(config, {"checkpoint_id": parent_id}), + checkpoint.checkpoint, + checkpoint.metadata, + checkpoint.checkpoint["channel_versions"], ) + parent_id = checkpoint.checkpoint["id"] + + # load history + history = [c async for c in app.aget_state_history(config)] + # check history with saved checkpoints matches expected history + exc_task_results: int = 0 + if checkpoint_version == "2-start:*": + exc_task_results = 1 + elif checkpoint_version == "2-quadratic": + exc_task_results = 2 + expected_history = get_expected_history(exc_task_results=exc_task_results) + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + # resume from 2nd to latest checkpoint + assert [c async for c in app.astream(Command(resume=""), history[1].config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + # new checkpoint should match the latest checkpoint in history + latest_state = await app.aget_state(config) + assert ( + StateSnapshot( + values=latest_state.values, + next=latest_state.next, + config=patch_configurable(latest_state.config, {"checkpoint_id": AnyStr()}), + metadata=AnyDict(latest_state.metadata), + created_at=AnyStr(), + parent_config=latest_state.parent_config, + tasks=latest_state.tasks, + interrupts=latest_state.interrupts, + ) + == history[0] + ) diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index f305d94d6..6b86129fc 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -1,20 +1,14 @@ import pytest from typing_extensions import TypedDict +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.graph import END, START, StateGraph -from tests.conftest import ( - REGULAR_CHECKPOINTERS_ASYNC, - REGULAR_CHECKPOINTERS_SYNC, - awith_checkpointer, -) pytestmark = pytest.mark.anyio -@pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -34,8 +28,7 @@ def test_interruption_without_state_updates( builder.add_edge("step_2", "step_3") builder.add_edge("step_3", END) - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - graph = builder.compile(checkpointer=checkpointer, interrupt_after="*") + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_after="*") initial_input = {"input": "hello world"} thread = {"configurable": {"thread_id": "1"}} @@ -56,10 +49,8 @@ def test_interruption_without_state_updates( assert n_checkpoints == (5 if checkpoint_during else 3) -@pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_interruption_without_state_updates_async( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -79,23 +70,22 @@ async def test_interruption_without_state_updates_async( builder.add_edge("step_2", "step_3") builder.add_edge("step_3", END) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_after="*") + graph = builder.compile(checkpointer=async_checkpointer, interrupt_after="*") - initial_input = {"input": "hello world"} - thread = {"configurable": {"thread_id": "1"}} + initial_input = {"input": "hello world"} + thread = {"configurable": {"thread_id": "1"}} - await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during) - assert (await graph.aget_state(thread)).next == ("step_2",) - n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) - assert n_checkpoints == (3 if checkpoint_during else 1) + await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during) + assert (await graph.aget_state(thread)).next == ("step_2",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (3 if checkpoint_during else 1) - await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) - assert (await graph.aget_state(thread)).next == ("step_3",) - n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) - assert n_checkpoints == (4 if checkpoint_during else 2) + await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) + assert (await graph.aget_state(thread)).next == ("step_3",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (4 if checkpoint_during else 2) - await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) - assert (await graph.aget_state(thread)).next == () - n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) - assert n_checkpoints == (5 if checkpoint_during else 3) + await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during) + assert (await graph.aget_state(thread)).next == () + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (5 if checkpoint_during else 3) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 57827f143..eb73fdef3 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -41,11 +41,7 @@ from langgraph.types import ( from tests.agents import AgentAction, AgentFinish from tests.any_int import AnyInt from tests.any_str import AnyDict, AnyStr, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_SYNC, - REGULAR_CHECKPOINTERS_SYNC, - SHOULD_CHECK_SNAPSHOTS, -) +from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS from tests.fake_chat import FakeChatModel from tests.fake_tracer import FakeTracer from tests.messages import ( @@ -314,18 +310,16 @@ def test_invoke_two_processes_in_out_interrupt( ] -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_fork_always_re_runs_nodes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture ) -> None: - 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) + graph = builder.compile(checkpointer=sync_checkpointer) thread1 = {"configurable": {"thread_id": "1"}} @@ -7341,13 +7335,9 @@ def test_branch_then( ) -@pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class InterruptOnce: ticks: int = 0 @@ -7398,7 +7388,7 @@ def test_send_dedupe_on_resume( builder.add_conditional_edges("1", send_for_fun) builder.add_conditional_edges("2", route_to_three) - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) thread1 = {"configurable": {"thread_id": "1"}} assert graph.invoke(["0"], thread1, checkpoint_during=checkpoint_during) == { "__interrupt__": [ @@ -7413,8 +7403,6 @@ def test_send_dedupe_on_resume( assert builder.nodes["flaky"].runnable.func.ticks == 1 # check state state = graph.get_state(thread1) - if "shallow" in checkpointer_name: - pytest.xfail("TODO: shallow checkpointer reports wrong next set") assert state.next == ("flaky",) # check history history = [c for c in graph.get_state_history(thread1)] diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index c81b81e0f..4007636a7 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -25,6 +25,7 @@ from typing_extensions import TypedDict from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.untracked_value import UntrackedValue +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import END, PULL, PUSH, START from langgraph.graph.graph import Graph from langgraph.graph.message import MessageGraph, add_messages @@ -39,7 +40,6 @@ from tests.any_int import AnyInt from tests.any_str import AnyDict, AnyStr, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, - REGULAR_CHECKPOINTERS_ASYNC, awith_checkpointer, ) from tests.fake_chat import FakeChatModel @@ -327,9 +327,8 @@ async def test_invoke_two_processes_in_out_interrupt( ] -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_fork_always_re_runs_nodes( - checkpointer_name: str, mocker: MockerFixture + async_checkpointer: BaseCheckpointSaver, mocker: MockerFixture ) -> None: add_one = mocker.Mock(side_effect=lambda _: 1) @@ -337,208 +336,201 @@ async def test_fork_always_re_runs_nodes( 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) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_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=(), - tasks=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"add_one": 1}, + # list history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + tasks=(), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - interrupts=(), - ), - StateSnapshot( - values=5, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": {"add_one": 1}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[1].config, + interrupts=(), + ), + StateSnapshot( + values=5, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - interrupts=(), - ), - StateSnapshot( - values=4, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"add_one": 1}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[2].config, + interrupts=(), + ), + StateSnapshot( + values=4, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - interrupts=(), - ), - StateSnapshot( - values=3, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"add_one": 1}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[3].config, + interrupts=(), + ), + StateSnapshot( + values=3, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - interrupts=(), - ), - StateSnapshot( - values=2, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"add_one": 1}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[4].config, + interrupts=(), + ), + StateSnapshot( + values=2, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - interrupts=(), - ), - StateSnapshot( - values=1, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[5].config, + interrupts=(), + ), + StateSnapshot( + values=1, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - interrupts=(), - ), - StateSnapshot( - values=0, - tasks=( - PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": 1}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[6].config, + interrupts=(), + ), + StateSnapshot( + values=0, + tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),), + next=("__start__",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - ), - ] + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] - # forking from any previous checkpoint should re-run nodes - 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") - ] == [ - {"add_one": 1}, - ] - assert [ - c - async for c in graph.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"add_one": 1}, - {"add_one": 1}, - ] + # forking from any previous checkpoint should re-run nodes + 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") + ] == [ + {"add_one": 1}, + ] + assert [ + c async for c in graph.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 389f11427..18eb6cfd5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -69,11 +69,7 @@ from langgraph.types import ( ) from tests.agents import AgentAction, AgentFinish from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_SYNC, - REGULAR_CHECKPOINTERS_SYNC, - SHOULD_CHECK_SNAPSHOTS, -) +from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS from tests.messages import ( _AnyIdAIMessage, _AnyIdAIMessageChunk, @@ -675,12 +671,9 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_run_from_checkpoint_id_retains_previous_writes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class MyState(TypedDict): myval: Annotated[int, operator.add] otherval: bool @@ -713,7 +706,7 @@ def test_run_from_checkpoint_id_retains_previous_writes( builder.add_conditional_edges("node_one", _getedge("node_one")) builder.add_conditional_edges("node_two", _getedge("node_two")) - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) thread_id = uuid.uuid4() thread1 = {"configurable": {"thread_id": str(thread_id)}} @@ -1105,7 +1098,6 @@ def test_invoke_checkpoint_two( assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -1504,7 +1496,6 @@ def test_send_sequences() -> None: ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_task( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -1602,7 +1593,6 @@ def test_imp_task( assert mapper_calls == 2 -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_nested( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -1676,7 +1666,6 @@ def test_imp_nested( ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_stream_order( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -3978,7 +3967,6 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -4047,7 +4035,6 @@ def test_subgraph_checkpoint_true( ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true_interrupt( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -4229,7 +4216,6 @@ def test_stream_buffering_single_node( ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -4397,7 +4383,6 @@ def test_nested_graph_interrupts_parallel( ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool @@ -4741,9 +4726,8 @@ def test_checkpoint_metadata() -> None: assert chkpnt_tuple.metadata["test_config_4"] == "bar" -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_remove_message_via_state_update( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage @@ -4760,8 +4744,7 @@ def test_remove_message_via_state_update( workflow.set_entry_point("chatbot") workflow.add_edge("chatbot", END) - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - app = workflow.compile(checkpointer=checkpointer) + app = workflow.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} output = app.invoke([HumanMessage(content="Hi")], config=config) app.update_state(config, values=[RemoveMessage(id=output[-1].id)]) @@ -6202,10 +6185,8 @@ def test_concurrent_execution_thread_safety(): assert result["counter"] == 1 -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) -def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: str): +def test_checkpoint_recovery(sync_checkpointer: BaseCheckpointSaver): """Test recovery from checkpoints after failures.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class State(TypedDict): steps: Annotated[list[str], operator.add] @@ -6226,7 +6207,7 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} # First attempt should fail @@ -6244,9 +6225,6 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: result = graph.invoke({"steps": [], "attempt": 2}, config) assert result == {"steps": ["start", "node1", "node2"], "attempt": 2} - if "shallow" in checkpointer_name: - return - # Verify checkpoint history shows both attempts history = list(graph.get_state_history(config)) assert len(history) == 6 # Initial + failed attempt + successful attempt @@ -8345,12 +8323,9 @@ def test_pregel_loop_refcount(): gc.enable() -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_bulk_state_updates( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): foo: str baz: str @@ -8367,7 +8342,7 @@ def test_bulk_state_updates( .add_node("node_b", node_b) .add_edge(START, "node_a") .add_edge("node_a", "node_b") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} @@ -8397,7 +8372,7 @@ def test_bulk_state_updates( assert state.values == {"foo": "updated", "baz": "new"} # Check if there are only two checkpoints - checkpoints = list(checkpointer.list(config)) + checkpoints = list(sync_checkpointer.list(config)) assert len(checkpoints) == 2 assert checkpoints[0].metadata["writes"] == { "node_a": {"foo": "updated"}, @@ -8424,7 +8399,7 @@ def test_bulk_state_updates( state = graph.get_state(config) assert state.values == {"foo": "updated", "baz": "new"} - checkpoints = list(checkpointer.list(config)) + checkpoints = list(sync_checkpointer.list(config)) assert len(checkpoints) == 2 assert checkpoints[0].metadata["writes"] == { "node_a": {"foo": "updated"}, @@ -8489,12 +8464,7 @@ def test_pregel_node_copy() -> None: graph.nodes["agent"].copy({}) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) -def test_update_as_input( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_update_as_input(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): foo: str @@ -8510,7 +8480,7 @@ def test_update_as_input( .add_node("tool", tool) .add_edge(START, "agent") .add_edge("agent", "tool") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { @@ -8560,12 +8530,9 @@ def test_update_as_input( assert new_history == history -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) def test_batch_update_as_input( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): foo: str tasks: Annotated[list[int], operator.add] @@ -8593,7 +8560,7 @@ def test_batch_update_as_input( .add_node("task", task) .add_edge(START, "agent") .add_edge("agent", "map") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e01ee06fe..4ef1dd049 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -38,6 +38,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.base import ( + BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, @@ -68,8 +69,6 @@ from langgraph.types import ( from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, - ALL_CHECKPOINTERS_ASYNC_PLUS_NONE, - REGULAR_CHECKPOINTERS_ASYNC, SHOULD_CHECK_SNAPSHOTS, awith_checkpointer, ) @@ -1220,8 +1219,7 @@ async def test_step_timeout_on_stream_hang(stream_hang_s: float) -> None: assert inner_task_cancelled -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) -async def test_cancel_graph_astream(checkpointer_name: str) -> None: +async def test_cancel_graph_astream(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): value: Annotated[int, operator.add] @@ -1255,45 +1253,44 @@ async def test_cancel_graph_astream(checkpointer_name: str) -> None: builder.add_edge(START, "aparallelwhile") builder.add_edge("alittlewhile", "awhile") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_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 - # but we should have applied pending writes - if checkpointer is not None: - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 3} # 1 + 2 - assert state.next == ("aparallelwhile",) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - } + # checkpoint with output of "alittlewhile" should not be saved + # but we should have applied pending writes + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 3} # 1 + 2 + assert state.next == ("aparallelwhile",) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) -async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) -> None: +async def test_cancel_graph_astream_events_v2( + async_checkpointer: BaseCheckpointSaver, +) -> None: class State(TypedDict): value: int @@ -1327,46 +1324,44 @@ async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) builder.add_edge("alittlewhile", "awhile") builder.add_edge("awhile", "anotherwhile") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_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}} - await asyncio.sleep(0.1) - 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}} + await asyncio.sleep(0.1) + 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 == { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"alittlewhile": {"value": 2}}, - "thread_id": "2", - } + # checkpoint with output of "alittlewhile" should not be saved + state = await graph.aget_state(thread2) + assert state is not None + assert state.values == {"value": 2} + assert state.next == ("awhile",) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"alittlewhile": {"value": 2}}, + "thread_id": "2", + } async def test_node_schemas_custom_output() -> None: @@ -2003,7 +1998,6 @@ async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( checkpointer_name: str, checkpoint_during: bool @@ -2284,9 +2278,8 @@ async def test_pending_writes_resume( ) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_run_from_checkpoint_id_retains_previous_writes( - checkpointer_name: str, + async_checkpointer: BaseCheckpointSaver, ) -> None: class MyState(TypedDict): myval: Annotated[int, operator.add] @@ -2320,47 +2313,46 @@ async def test_run_from_checkpoint_id_retains_previous_writes( builder.add_conditional_edges("node_one", _getedge("node_one")) builder.add_conditional_edges("node_two", _getedge("node_two")) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_checkpointer) - thread_id = uuid.uuid4() - thread1 = {"configurable": {"thread_id": str(thread_id)}} + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} - result = await graph.ainvoke({"myval": 1}, thread1) - assert result["myval"] == 4 - history = [c async for c in graph.aget_state_history(thread1)] + result = await graph.ainvoke({"myval": 1}, thread1) + assert result["myval"] == 4 + history = [c async for c in graph.aget_state_history(thread1)] - assert len(history) == 4 - assert history[0].values == {"myval": 4, "otherval": False} - assert history[-1].values == {"myval": 0} + assert len(history) == 4 + assert history[0].values == {"myval": 4, "otherval": False} + assert history[-1].values == {"myval": 0} - second_run_config = { - **thread1, - "configurable": { - **thread1["configurable"], - "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], - }, - } - second_result = await graph.ainvoke(None, second_run_config) - assert second_result == {"myval": 5, "otherval": True} + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = await graph.ainvoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} - new_history = [ - c - async for c in graph.aget_state_history( - {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} - ) - ] + new_history = [ + c + async for c in graph.aget_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] - assert len(new_history) == len(history) + 1 - for original, new in zip(history, new_history[1:]): - assert original.values == new.values - assert original.next == new.next - assert original.metadata["step"] == new.metadata["step"] + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] - def _get_tasks(hist: list, start: int): - return [h.tasks for h in hist[start:]] + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] - assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) async def test_cond_edge_after_send() -> None: @@ -2509,7 +2501,6 @@ async def test_send_sequences(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None: if not checkpoint_during and "shallow" in checkpointer_name: @@ -2573,7 +2564,6 @@ async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> None: if not checkpoint_during and "shallow" in checkpointer_name: @@ -2649,7 +2639,6 @@ async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> No @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) -> None: if not checkpoint_during and "shallow" in checkpointer_name: @@ -2711,7 +2700,6 @@ async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_sync_from_async( checkpointer_name: str, checkpoint_during: bool @@ -2755,7 +2743,6 @@ async def test_imp_sync_from_async( @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_stream_order( checkpointer_name: str, checkpoint_during: bool @@ -2799,10 +2786,8 @@ async def test_imp_stream_order( ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: class InterruptOnce: ticks: int = 0 @@ -2851,329 +2836,324 @@ async def test_send_dedupe_on_resume( builder.add_conditional_edges("1", send_for_fun) builder.add_conditional_edges("2", route_to_three) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke( - ["0"], thread1, checkpoint_during=checkpoint_during - ) == { - "__interrupt__": [ - Interrupt( - value="Bahh", - resumable=False, - ns=None, - ), + graph = builder.compile(checkpointer=async_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1, checkpoint_during=checkpoint_during) == { + "__interrupt__": [ + Interrupt( + value="Bahh", + resumable=False, + ns=None, + ), + ], + } + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # resume execution + assert await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + "3", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check history + history = [c async for c in graph.aget_state_history(thread1)] + assert len(history) == (6 if checkpoint_during else 2) + expected_history = [ + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + "3", ], - } - assert builder.nodes["2"].runnable.func.ticks == 3 - assert builder.nodes["flaky"].runnable.func.ticks == 1 - # resume execution - assert await graph.ainvoke( - None, thread1, checkpoint_during=checkpoint_during - ) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - "flaky|4", - "3", - ] - # node "2" doesn't get called again, as we recover writes saved before - assert builder.nodes["2"].runnable.func.ticks == 3 - # node "flaky" gets called again, as it was interrupted - assert builder.nodes["flaky"].runnable.func.ticks == 2 - # check history - history = [c async for c in graph.aget_state_history(thread1)] - assert len(history) == (6 if checkpoint_during else 2) - expected_history = [ - StateSnapshot( - values=[ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - "flaky|4", - "3", - ], - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"3": ["3"]}, + next=(), + config={ + "configurable": { "thread_id": "1", - "step": 4, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - interrupts=(), - ), - StateSnapshot( - values=[ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "3", - "2|3", - "flaky|4", - ], - next=("3",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"3": ["3"]}, + "thread_id": "1", + "step": 4, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { "thread_id": "1", - "step": 3, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="3", - path=("__pregel_pull", "3"), - error=None, - interrupts=(), - state=None, - result=["3"], - ), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + interrupts=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + ], + next=("3",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]}, + "thread_id": "1", + "step": 3, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], ), - interrupts=(), ), - StateSnapshot( - values=[ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - ], - next=("2", "flaky", "3"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "2": [ - ["2|Command(goto=Send(node='2', arg=3))"], - ["2|Command(goto=Send(node='flaky', arg=4))"], - ], - "3.1": ["3.1"], - }, + interrupts=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + ], + next=("2", "flaky", "3"), + config={ + "configurable": { "thread_id": "1", - "step": 2, - "parents": {}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "2": [ + ["2|Command(goto=Send(node='2', arg=3))"], + ["2|Command(goto=Send(node='flaky', arg=4))"], + ], + "3.1": ["3.1"], }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="2", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result=["2|3"], - ), - PregelTask( - id=AnyStr(), - name="flaky", - path=("__pregel_push", 1, False), - error=None, - interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), - state=None, - result=["flaky|4"] if checkpoint_during else None, - ), - PregelTask( - id=AnyStr(), - name="3", - path=("__pregel_pull", "3"), - error=None, - interrupts=(), - state=None, - result=["3"], - ), - ), - interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), - ), - StateSnapshot( - values=["0", "1"], - next=("2", "2", "3.1"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"1": ["1"]}, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { "thread_id": "1", - "step": 1, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="2", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='2', arg=3))"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=("__pregel_push", 1, False), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='flaky', arg=4))"], - ), - PregelTask( - id=AnyStr(), - name="3.1", - path=("__pregel_pull", "3.1"), - error=None, - interrupts=(), - state=None, - result=["3.1"], - ), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=("__pregel_push", 1, False), + error=None, + interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), + state=None, + result=["flaky|4"] if checkpoint_during else None, + ), + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], ), - interrupts=(), ), - StateSnapshot( - values=["0"], - next=("1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": None, + interrupts=(Interrupt(value="Bahh", resumable=False, ns=None),), + ), + StateSnapshot( + values=["0", "1"], + next=("2", "2", "3.1"), + config={ + "configurable": { "thread_id": "1", - "step": 0, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="1", - path=("__pregel_pull", "1"), - error=None, - interrupts=(), - state=None, - result=["1"], - ), - ), - interrupts=(), - ), - StateSnapshot( - values=[], - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": ["0"]}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"1": ["1"]}, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { "thread_id": "1", - "step": -1, - "parents": {}, - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result=["0"], - ), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 1, False), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], ), - interrupts=(), ), - ] - if checkpoint_during: - assert history == expected_history - else: - assert history[0] == expected_history[0] - assert history[1] == expected_history[2] + interrupts=(), + ), + StateSnapshot( + values=["0"], + next=("1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": ["0"]}, + "thread_id": "1", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + interrupts=(), + ), + ] + if checkpoint_during: + assert history == expected_history + else: + assert history[0] == expected_history[0] + assert history[1] == expected_history[2] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) @@ -5724,10 +5704,8 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 -@pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_subgraph_checkpoint_true( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] @@ -5756,48 +5734,45 @@ async def test_subgraph_checkpoint_true( "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END ) - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=async_checkpointer) - config = {"configurable": {"thread_id": "2"}} - assert [ - c - async for c in app.astream( - {"my_key": ""}, - config, - subgraphs=True, - checkpoint_during=checkpoint_during, - ) - ] == [ - (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), - (("inner",), {"inner_2": {"my_key": " and there"}}), - ((), {"inner": {"my_key": " got here and there"}}), - ( - ("inner",), - { - "inner_1": { - "my_key": " got here", - "my_other_key": " got here and there got here and there", - } - }, - ), - (("inner",), {"inner_2": {"my_key": " and there"}}), - ( - (), - { - "inner": { - "my_key": " got here and there got here and there got here and there" - } - }, - ), - ] + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) + ] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpoint_during", [True, False]) -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_subgraph_checkpoint_true_interrupt( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: # Define subgraph class SubgraphState(TypedDict): @@ -5835,28 +5810,27 @@ async def test_subgraph_checkpoint_true_interrupt( builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke( - {"foo": "foo"}, config, checkpoint_during=checkpoint_during - ) == { - "foo": "hi! foo", - "__interrupt__": [ - Interrupt( - value="Provide baz value", - resumable=True, - ns=[AnyStr("node_2"), AnyStr("subgraph_node_1:")], - ) - ], - } - assert (await graph.aget_state(config, subgraphs=True)).tasks[ - 0 - ].state.values == {"bar": "hi! foo"} - assert await graph.ainvoke( - Command(resume="baz"), config, checkpoint_during=checkpoint_during - ) == {"foo": "hi! foobaz"} + assert await graph.ainvoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == { + "foo": "hi! foo", + "__interrupt__": [ + Interrupt( + value="Provide baz value", + resumable=True, + ns=[AnyStr("node_2"), AnyStr("subgraph_node_1:")], + ) + ], + } + assert (await graph.aget_state(config, subgraphs=True)).tasks[0].state.values == { + "bar": "hi! foo" + } + assert await graph.ainvoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) @@ -5967,7 +5941,6 @@ async def test_stream_buffering_single_node(checkpointer_name: str) -> None: ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_interrupts_parallel( checkpointer_name: str, checkpoint_during: bool @@ -6158,7 +6131,6 @@ async def test_nested_graph_interrupts_parallel( ] -@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_doubly_nested_graph_interrupts( checkpointer_name: str, checkpoint_during: bool @@ -8705,292 +8677,275 @@ async def test_pregel_loop_refcount(): gc.enable() -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_bulk_state_updates(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: +async def test_bulk_state_updates(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + foo: str + baz: str - class State(TypedDict): - foo: str - baz: str + def node_a(state: State) -> State: + return {"foo": "bar"} - def node_a(state: State) -> State: - return {"foo": "bar"} + def node_b(state: State) -> State: + return {"baz": "qux"} - def node_b(state: State) -> State: - return {"baz": "qux"} + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=async_checkpointer) + ) - graph = ( - StateGraph(State) - .add_node("node_a", node_a) - .add_node("node_b", node_b) - .add_edge(START, "node_a") - .add_edge("node_a", "node_b") - .compile(checkpointer=checkpointer) - ) + config = {"configurable": {"thread_id": "1"}} - config = {"configurable": {"thread_id": "1"}} + # First update with node_a + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "bar"}, "node_a"), + ] + ], + ) - # First update with node_a + # Then bulk update with both nodes + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "updated"}, "node_a"), + StateUpdate({"baz": "new"}, "node_b"), + ] + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = [ + c async for c in async_checkpointer.alist({"configurable": {"thread_id": "1"}}) + ] + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # perform multiple steps at the same time + config = {"configurable": {"thread_id": "2"}} + + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "bar"}, "node_a"), + ], + [ + StateUpdate({"foo": "updated"}, "node_a"), + StateUpdate({"baz": "new"}, "node_b"), + ], + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + checkpoints = [ + c async for c in async_checkpointer.alist({"configurable": {"thread_id": "1"}}) + ] + assert len(checkpoints) == 2 + assert checkpoints[0].metadata["writes"] == { + "node_a": {"foo": "updated"}, + "node_b": {"baz": "new"}, + } + assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): await graph.abulk_update_state( config, [ [ - StateUpdate({"foo": "bar"}, "node_a"), + StateUpdate(values={"foo": "error"}, as_node=None), + StateUpdate(values={"bar": "error"}, as_node=None), ] ], ) - # Then bulk update with both nodes + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No supersteps provided"): + await graph.abulk_update_state(config, []) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + await graph.abulk_update_state(config, [[], []]) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): await graph.abulk_update_state( config, [ [ - StateUpdate({"foo": "updated"}, "node_a"), - StateUpdate({"baz": "new"}, "node_b"), - ] + StateUpdate(values=None, as_node="__end__"), + StateUpdate(values=None, as_node="__copy__"), + ], ], ) - state = await graph.aget_state(config) - assert state.values == {"foo": "updated", "baz": "new"} - # Check if there are only two checkpoints - checkpoints = [ - c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}}) - ] - assert len(checkpoints) == 2 - assert checkpoints[0].metadata["writes"] == { - "node_a": {"foo": "updated"}, - "node_b": {"baz": "new"}, +async def test_update_as_input(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=async_checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "tool"} + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "tool"} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), } - assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} - # perform multiple steps at the same time - config = {"configurable": {"thread_id": "2"}} + history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "1"}}) + ] - await graph.abulk_update_state( - config, - [ - [ - StateUpdate({"foo": "bar"}, "node_a"), - ], - [ - StateUpdate({"foo": "updated"}, "node_a"), - StateUpdate({"baz": "new"}, "node_b"), - ], + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}}) + ] + + assert new_history == history + + +async def test_batch_update_as_input(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), ], + update={"foo": "map"}, ) - state = await graph.aget_state(config) - assert state.values == {"foo": "updated", "baz": "new"} + def task(state: dict) -> State: + return {"tasks": [state["index"]]} - checkpoints = [ - c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}}) - ] - assert len(checkpoints) == 2 - assert checkpoints[0].metadata["writes"] == { - "node_a": {"foo": "updated"}, - "node_b": {"baz": "new"}, + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=async_checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, {"configurable": {"thread_id": "1"}} + ) == {"foo": "map", "tasks": [0, 1, 2]} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], } - assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}} - # Should raise error if updating without as_node - with pytest.raises(InvalidUpdateError): - await graph.abulk_update_state( - config, - [ - [ - StateUpdate(values={"foo": "error"}, as_node=None), - StateUpdate(values={"bar": "error"}, as_node=None), - ] - ], - ) + history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "1"}}) + ] - # Should raise if no updates are provided - with pytest.raises(ValueError, match="No supersteps provided"): - await graph.abulk_update_state(config, []) - - # Should raise if no updates are provided - with pytest.raises(ValueError, match="No updates provided"): - await graph.abulk_update_state(config, [[], []]) - - # Should raise if __end__ or __copy__ update is applied in bulk - with pytest.raises(InvalidUpdateError): - await graph.abulk_update_state( - config, - [ - [ - StateUpdate(values=None, as_node="__end__"), - StateUpdate(values=None, as_node="__copy__"), - ], - ], - ) - - -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_update_as_input(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - - class State(TypedDict): - foo: str - - def agent(state: State) -> State: - return {"foo": "agent"} - - def tool(state: State) -> State: - return {"foo": "tool"} - - graph = ( - StateGraph(State) - .add_node("agent", agent) - .add_node("tool", tool) - .add_edge(START, "agent") - .add_edge("agent", "tool") - .compile(checkpointer=checkpointer) - ) - - assert await graph.ainvoke( - {"foo": "input"}, {"configurable": {"thread_id": "1"}} - ) == {"foo": "tool"} - - assert await graph.ainvoke( - {"foo": "input"}, {"configurable": {"thread_id": "1"}} - ) == {"foo": "tool"} - - def map_snapshot(i: StateSnapshot) -> dict: - return { - "values": i.values, - "next": i.next, - "step": i.metadata.get("step"), - } - - history = [ - map_snapshot(s) - async for s in graph.aget_state_history( - {"configurable": {"thread_id": "1"}} - ) - ] - - await graph.abulk_update_state( - {"configurable": {"thread_id": "2"}}, + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], [ - # First turn - [StateUpdate({"foo": "input"}, "__input__")], - [StateUpdate({"foo": "input"}, "__start__")], - [StateUpdate({"foo": "agent"}, "agent")], - [StateUpdate({"foo": "tool"}, "tool")], - # Second turn - [StateUpdate({"foo": "input"}, "__input__")], - [StateUpdate({"foo": "input"}, "__start__")], - [StateUpdate({"foo": "agent"}, "agent")], - [StateUpdate({"foo": "tool"}, "tool")], + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) ], - ) - - state = await graph.aget_state({"configurable": {"thread_id": "2"}}) - assert state.values == {"foo": "tool"} - - new_history = [ - map_snapshot(s) - async for s in graph.aget_state_history( - {"configurable": {"thread_id": "2"}} - ) - ] - - assert new_history == history - - -@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) -async def test_batch_update_as_input(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - - class State(TypedDict): - foo: str - tasks: Annotated[list[int], operator.add] - - def agent(state: State) -> State: - return {"foo": "agent"} - - def map(state: State) -> Command["task"]: - return Command( - goto=[ - Send("task", {"index": 0}), - Send("task", {"index": 1}), - Send("task", {"index": 2}), - ], - update={"foo": "map"}, - ) - - def task(state: dict) -> State: - return {"tasks": [state["index"]]} - - graph = ( - StateGraph(State) - .add_node("agent", agent) - .add_node("map", map) - .add_node("task", task) - .add_edge(START, "agent") - .add_edge("agent", "map") - .compile(checkpointer=checkpointer) - ) - - assert await graph.ainvoke( - {"foo": "input"}, {"configurable": {"thread_id": "1"}} - ) == {"foo": "map", "tasks": [0, 1, 2]} - - def map_snapshot(i: StateSnapshot) -> dict: - return { - "values": i.values, - "next": i.next, - "step": i.metadata.get("step"), - "tasks": [t.name for t in i.tasks], - } - - history = [ - map_snapshot(s) - async for s in graph.aget_state_history( - {"configurable": {"thread_id": "1"}} - ) - ] - - await graph.abulk_update_state( - {"configurable": {"thread_id": "2"}}, [ - [StateUpdate({"foo": "input"}, "__input__")], - [StateUpdate({"foo": "input"}, "__start__")], - [StateUpdate({"foo": "agent", "tasks": []}, "agent")], - [ - StateUpdate( - Command( - goto=[ - Send("task", {"index": 0}), - Send("task", {"index": 1}), - Send("task", {"index": 2}), - ], - update={"foo": "map"}, - ), - "map", - ) - ], - [ - StateUpdate({"tasks": [0]}, "task"), - StateUpdate({"tasks": [1]}, "task"), - StateUpdate({"tasks": [2]}, "task"), - ], + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), ], - ) + ], + ) - state = await graph.aget_state({"configurable": {"thread_id": "2"}}) - assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} - new_history = [ - map_snapshot(s) - async for s in graph.aget_state_history( - {"configurable": {"thread_id": "2"}} - ) - ] + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}}) + ] - assert new_history == history + assert new_history == history async def test_draw_invalid(): From ed7f038a190f67b7304b3d1a609b249caae56f91 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 10:54:12 -0700 Subject: [PATCH 2/2] Fix --- libs/langgraph/tests/test_checkpoint_migration.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/tests/test_checkpoint_migration.py b/libs/langgraph/tests/test_checkpoint_migration.py index caba88060..f6583cada 100644 --- a/libs/langgraph/tests/test_checkpoint_migration.py +++ b/libs/langgraph/tests/test_checkpoint_migration.py @@ -1672,16 +1672,11 @@ async def test_latest_checkpoint_state_graph_async( @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"]) def test_saved_checkpoint_state_graph( - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, checkpoint_version: str, ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - builder = make_state_graph() - app = builder.compile(checkpointer=checkpointer) + app = builder.compile(checkpointer=sync_checkpointer) thread1 = "1" config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} @@ -1693,8 +1688,8 @@ def test_saved_checkpoint_state_graph( for write in checkpoint.pending_writes: grouped_writes[write[0]].append(write[1:]) for tid, group in grouped_writes.items(): - checkpointer.put_writes(checkpoint.config, group, tid) - checkpointer.put( + sync_checkpointer.put_writes(checkpoint.config, group, tid) + sync_checkpointer.put( patch_configurable(config, {"checkpoint_id": parent_id}), checkpoint.checkpoint, checkpoint.metadata,