From 8c11c1155ae6b02f1a6a564f4813e02082e6097c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 May 2025 12:39:50 -0700 Subject: [PATCH] Remove postgres shallow checkpointer - This was deprecated, and superseded by checkpoint_during=False, which is available for all checkpointers --- .../langgraph/checkpoint/postgres/__init__.py | 3 +- .../langgraph/checkpoint/postgres/aio.py | 3 +- .../langgraph/checkpoint/postgres/shallow.py | 941 ------------------ libs/checkpoint-postgres/tests/test_async.py | 41 +- libs/checkpoint-postgres/tests/test_sync.py | 34 +- libs/langgraph/.claude/settings.local.json | 9 + libs/langgraph/tests/conftest.py | 13 - libs/langgraph/tests/conftest_checkpointer.py | 53 +- libs/langgraph/tests/test_large_cases.py | 630 +++--------- .../langgraph/tests/test_large_cases_async.py | 429 +++----- libs/langgraph/tests/test_pregel.py | 41 +- libs/langgraph/tests/test_pregel_async.py | 193 ++-- 12 files changed, 366 insertions(+), 2024 deletions(-) delete mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py create mode 100644 libs/langgraph/.claude/settings.local.json diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 4da709495..3a2ecbce9 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -20,7 +20,6 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres.base import BasePostgresSaver -from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol Conn = _internal.Conn # For backward compatibility @@ -425,4 +424,4 @@ class PostgresSaver(BasePostgresSaver): yield cur -__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"] +__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index e750203de..ddbc2fdbc 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -20,7 +20,6 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.postgres import _ainternal from langgraph.checkpoint.postgres.base import BasePostgresSaver -from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol Conn = _ainternal.Conn # For backward compatibility @@ -532,4 +531,4 @@ class AsyncPostgresSaver(BasePostgresSaver): ).result() -__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"] +__all__ = ["AsyncPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py deleted file mode 100644 index 4677f4f3d..000000000 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py +++ /dev/null @@ -1,941 +0,0 @@ -import asyncio -import threading -import warnings -from collections.abc import AsyncIterator, Iterator, Sequence -from contextlib import asynccontextmanager, contextmanager -from typing import Any, Optional - -from langchain_core.runnables import RunnableConfig -from psycopg import ( - AsyncConnection, - AsyncCursor, - AsyncPipeline, - Capabilities, - Connection, - Cursor, - Pipeline, -) -from psycopg.rows import DictRow, dict_row -from psycopg.types.json import Jsonb -from psycopg_pool import AsyncConnectionPool, ConnectionPool - -from langgraph.checkpoint.base import ( - WRITES_IDX_MAP, - ChannelVersions, - Checkpoint, - CheckpointMetadata, - CheckpointTuple, - get_checkpoint_metadata, -) -from langgraph.checkpoint.postgres import _ainternal, _internal -from langgraph.checkpoint.postgres.base import BasePostgresSaver -from langgraph.checkpoint.serde.base import SerializerProtocol -from langgraph.checkpoint.serde.types import TASKS - -""" -To add a new migration, add a new string to the MIGRATIONS list. -The position of the migration in the list is the version number. -""" -MIGRATIONS = [ - """CREATE TABLE IF NOT EXISTS checkpoint_migrations ( - v INTEGER PRIMARY KEY -);""", - """CREATE TABLE IF NOT EXISTS checkpoints ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - type TEXT, - checkpoint JSONB NOT NULL, - metadata JSONB NOT NULL DEFAULT '{}', - PRIMARY KEY (thread_id, checkpoint_ns) -);""", - """CREATE TABLE IF NOT EXISTS checkpoint_blobs ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - channel TEXT NOT NULL, - type TEXT NOT NULL, - blob BYTEA, - PRIMARY KEY (thread_id, checkpoint_ns, channel) -);""", - """CREATE TABLE IF NOT EXISTS checkpoint_writes ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - checkpoint_id TEXT NOT NULL, - task_id TEXT NOT NULL, - idx INTEGER NOT NULL, - channel TEXT NOT NULL, - type TEXT, - blob BYTEA NOT NULL, - PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) -);""", - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id); - """, - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id); - """, - """ - CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id); - """, - """ - ALTER TABLE checkpoint_writes ADD COLUMN task_path TEXT NOT NULL DEFAULT ''; - """, -] - -SELECT_SQL = f""" -select - thread_id, - checkpoint, - checkpoint_ns, - metadata, - ( - select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob]) - from jsonb_each_text(checkpoint -> 'channel_versions') - inner join checkpoint_blobs bl - on bl.thread_id = checkpoints.thread_id - and bl.checkpoint_ns = checkpoints.checkpoint_ns - and bl.channel = jsonb_each_text.key - ) as channel_values, - ( - select - array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob] order by cw.task_id, cw.idx) - from checkpoint_writes cw - where cw.thread_id = checkpoints.thread_id - and cw.checkpoint_ns = checkpoints.checkpoint_ns - and cw.checkpoint_id = (checkpoint->>'id') - ) as pending_writes, - ( - select array_agg(array[cw.type::bytea, cw.blob] order by cw.task_path, cw.task_id, cw.idx) - from checkpoint_writes cw - where cw.thread_id = checkpoints.thread_id - and cw.checkpoint_ns = checkpoints.checkpoint_ns - and cw.channel = '{TASKS}' - ) as pending_sends -from checkpoints """ - -UPSERT_CHECKPOINT_BLOBS_SQL = """ - INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, type, blob) - VALUES (%s, %s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns, channel) DO UPDATE SET - type = EXCLUDED.type, - blob = EXCLUDED.blob; -""" - -UPSERT_CHECKPOINTS_SQL = """ - INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint, metadata) - VALUES (%s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns) - DO UPDATE SET - checkpoint = EXCLUDED.checkpoint, - metadata = EXCLUDED.metadata; -""" - -UPSERT_CHECKPOINT_WRITES_SQL = """ - INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET - channel = EXCLUDED.channel, - type = EXCLUDED.type, - blob = EXCLUDED.blob; -""" - -INSERT_CHECKPOINT_WRITES_SQL = """ - INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, task_path, idx, channel, type, blob) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING -""" - - -def _dump_blobs( - serde: SerializerProtocol, - thread_id: str, - checkpoint_ns: str, - values: dict[str, Any], - versions: ChannelVersions, -) -> list[tuple[str, str, str, str, Optional[bytes]]]: - if not versions: - return [] - - return [ - ( - thread_id, - checkpoint_ns, - k, - *(serde.dumps_typed(values[k]) if k in values else ("empty", None)), - ) - for k in versions - ] - - -class ShallowPostgresSaver(BasePostgresSaver): - """A checkpoint saver that uses Postgres to store checkpoints. - - This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history. - It is meant to be a light-weight drop-in replacement for the PostgresSaver that - supports most of the LangGraph persistence functionality with the exception of time travel. - """ - - SELECT_SQL = SELECT_SQL - MIGRATIONS = MIGRATIONS - UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL - UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL - UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL - INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL - - lock: threading.Lock - - def __init__( - self, - conn: _internal.Conn, - pipe: Optional[Pipeline] = None, - serde: Optional[SerializerProtocol] = None, - ) -> None: - warnings.warn( - "ShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. " - "Use PostgresSaver instead, and invoke the graph with `graph.invoke(..., checkpoint_during=False)`.", - DeprecationWarning, - stacklevel=2, - ) - super().__init__(serde=serde) - if isinstance(conn, ConnectionPool) and pipe is not None: - raise ValueError( - "Pipeline should be used only with a single Connection, not ConnectionPool." - ) - - self.conn = conn - self.pipe = pipe - self.lock = threading.Lock() - self.supports_pipeline = Capabilities().has_pipeline() - - @classmethod - @contextmanager - def from_conn_string( - cls, conn_string: str, *, pipeline: bool = False - ) -> Iterator["ShallowPostgresSaver"]: - """Create a new ShallowPostgresSaver instance from a connection string. - - Args: - conn_string: The Postgres connection info string. - pipeline: whether to use Pipeline - - Returns: - ShallowPostgresSaver: A new ShallowPostgresSaver instance. - """ - with Connection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - if pipeline: - with conn.pipeline() as pipe: - yield cls(conn, pipe) - else: - yield cls(conn) - - def setup(self) -> None: - """Set up the checkpoint database asynchronously. - - This method creates the necessary tables in the Postgres database if they don't - already exist and runs database migrations. It MUST be called directly by the user - the first time checkpointer is used. - """ - with self._cursor() as cur: - cur.execute(self.MIGRATIONS[0]) - results = cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ) - row = results.fetchone() - if row is None: - version = -1 - else: - version = row["v"] - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - cur.execute(migration) - cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") - if self.pipe: - self.pipe.sync() - - def list( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. For ShallowPostgresSaver, this method returns a list with - ONLY the most recent checkpoint. - """ - where, args = self._search_where(config, filter, before) - query = self.SELECT_SQL + where - if limit: - query += f" LIMIT {limit}" - with self._cursor() as cur: - cur.execute(self.SELECT_SQL + where, args, binary=True) - for value in cur: - checkpoint = self._load_checkpoint( - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=self._load_writes(value["pending_writes"]), - ) - - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config (matching the thread ID in the config). - - Args: - config: The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - - Examples: - - Basic: - >>> config = {"configurable": {"thread_id": "1"}} - >>> checkpoint_tuple = memory.get_tuple(config) - >>> print(checkpoint_tuple) - CheckpointTuple(...) - - With timestamp: - - >>> config = { - ... "configurable": { - ... "thread_id": "1", - ... "checkpoint_ns": "", - ... "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875", - ... } - ... } - >>> checkpoint_tuple = memory.get_tuple(config) - >>> print(checkpoint_tuple) - CheckpointTuple(...) - """ # noqa - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - args = (thread_id, checkpoint_ns) - where = "WHERE thread_id = %s AND checkpoint_ns = %s" - - with self._cursor() as cur: - cur.execute( - self.SELECT_SQL + where, - args, - binary=True, - ) - - for value in cur: - checkpoint = self._load_checkpoint( - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - return CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=self._load_writes(value["pending_writes"]), - ) - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config. For ShallowPostgresSaver, this method saves ONLY the most recent - checkpoint and overwrites a previous checkpoint, if it exists. - - Args: - config: The config to associate with the checkpoint. - checkpoint: The checkpoint to save. - metadata: Additional metadata to save with the checkpoint. - new_versions: New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - - Examples: - - >>> from langgraph.checkpoint.postgres import ShallowPostgresSaver - >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" - >>> with ShallowPostgresSaver.from_conn_string(DB_URI) as memory: - >>> config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} - >>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "channel_values": {"key": "value"}} - >>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {}) - >>> print(saved_config) - {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}} - """ - configurable = config["configurable"].copy() - thread_id = configurable.pop("thread_id") - checkpoint_ns = configurable.pop("checkpoint_ns") - - copy = checkpoint.copy() - next_config = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - } - - with self._cursor(pipeline=True) as cur: - cur.execute( - """DELETE FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""", - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - configurable.get("checkpoint_id", ""), - ), - ) - cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - _dump_blobs( - self.serde, - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(get_checkpoint_metadata(config, metadata)), - ), - ) - return next_config - - def put_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint. - - This method saves intermediate writes associated with a checkpoint to the Postgres database. - - Args: - config: Configuration of the related checkpoint. - writes: List of writes to store. - task_id: Identifier for the task creating the writes. - """ - query = ( - self.UPSERT_CHECKPOINT_WRITES_SQL - if all(w[0] in WRITES_IDX_MAP for w in writes) - else self.INSERT_CHECKPOINT_WRITES_SQL - ) - with self._cursor(pipeline=True) as cur: - cur.executemany( - query, - self._dump_writes( - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - task_path, - writes, - ), - ) - - @contextmanager - def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: - """Create a database cursor as a context manager. - - Args: - pipeline: whether to use pipeline for the DB operations inside the context manager. - Will be applied regardless of whether the ShallowPostgresSaver instance was initialized with a pipeline. - If pipeline mode is not supported, will fall back to using transaction context manager. - """ - with _internal.get_connection(self.conn) as conn: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - with conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - finally: - if pipeline: - self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - if self.supports_pipeline: - with ( - self.lock, - conn.pipeline(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - # Use connection's transaction context manager when pipeline mode not supported - with ( - self.lock, - conn.transaction(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - - -class AsyncShallowPostgresSaver(BasePostgresSaver): - """A checkpoint saver that uses Postgres to store checkpoints asynchronously. - - This checkpointer ONLY stores the most recent checkpoint and does NOT retain any history. - It is meant to be a light-weight drop-in replacement for the AsyncPostgresSaver that - supports most of the LangGraph persistence functionality with the exception of time travel. - """ - - SELECT_SQL = SELECT_SQL - MIGRATIONS = MIGRATIONS - UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL - UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL - UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL - INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL - lock: asyncio.Lock - - def __init__( - self, - conn: _ainternal.Conn, - pipe: Optional[AsyncPipeline] = None, - serde: Optional[SerializerProtocol] = None, - ) -> None: - warnings.warn( - "AsyncShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. " - "Use AsyncPostgresSaver instead, and invoke the graph with `await graph.ainvoke(..., checkpoint_during=False)`.", - DeprecationWarning, - stacklevel=2, - ) - super().__init__(serde=serde) - if isinstance(conn, AsyncConnectionPool) and pipe is not None: - raise ValueError( - "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." - ) - - self.conn = conn - self.pipe = pipe - self.lock = asyncio.Lock() - self.loop = asyncio.get_running_loop() - self.supports_pipeline = Capabilities().has_pipeline() - - @classmethod - @asynccontextmanager - async def from_conn_string( - cls, - conn_string: str, - *, - pipeline: bool = False, - serde: Optional[SerializerProtocol] = None, - ) -> AsyncIterator["AsyncShallowPostgresSaver"]: - """Create a new AsyncShallowPostgresSaver instance from a connection string. - - Args: - conn_string: The Postgres connection info string. - pipeline: whether to use AsyncPipeline - - Returns: - AsyncShallowPostgresSaver: A new AsyncShallowPostgresSaver instance. - """ - async with await AsyncConnection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - if pipeline: - async with conn.pipeline() as pipe: - yield cls(conn=conn, pipe=pipe, serde=serde) - else: - yield cls(conn=conn, serde=serde) - - async def setup(self) -> None: - """Set up the checkpoint database asynchronously. - - This method creates the necessary tables in the Postgres database if they don't - already exist and runs database migrations. It MUST be called directly by the user - the first time checkpointer is used. - """ - async with self._cursor() as cur: - await cur.execute(self.MIGRATIONS[0]) - results = await cur.execute( - "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" - ) - row = await results.fetchone() - if row is None: - version = -1 - else: - version = row["v"] - for v, migration in zip( - range(version + 1, len(self.MIGRATIONS)), - self.MIGRATIONS[version + 1 :], - ): - await cur.execute(migration) - await cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") - if self.pipe: - await self.pipe.sync() - - async def alist( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> AsyncIterator[CheckpointTuple]: - """List checkpoints from the database asynchronously. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. For ShallowPostgresSaver, this method returns a list with - ONLY the most recent checkpoint. - """ - where, args = self._search_where(config, filter, before) - query = self.SELECT_SQL + where - if limit: - query += f" LIMIT {limit}" - async with self._cursor() as cur: - await cur.execute(self.SELECT_SQL + where, args, binary=True) - async for value in cur: - checkpoint = await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=await asyncio.to_thread( - self._load_writes, value["pending_writes"] - ), - ) - - async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database asynchronously. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config (matching the thread ID in the config). - - Args: - config: The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - args = (thread_id, checkpoint_ns) - where = "WHERE thread_id = %s AND checkpoint_ns = %s" - - async with self._cursor() as cur: - await cur.execute( - self.SELECT_SQL + where, - args, - binary=True, - ) - - async for value in cur: - checkpoint = await asyncio.to_thread( - self._load_checkpoint, - value["checkpoint"], - value["channel_values"], - value["pending_sends"], - ) - return CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - }, - checkpoint=checkpoint, - metadata=self._load_metadata(value["metadata"]), - pending_writes=await asyncio.to_thread( - self._load_writes, value["pending_writes"] - ), - ) - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database asynchronously. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent - checkpoint and overwrites a previous checkpoint, if it exists. - - Args: - config: The config to associate with the checkpoint. - checkpoint: The checkpoint to save. - metadata: Additional metadata to save with the checkpoint. - new_versions: New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - """ - configurable = config["configurable"].copy() - thread_id = configurable.pop("thread_id") - checkpoint_ns = configurable.pop("checkpoint_ns") - - copy = checkpoint.copy() - next_config = { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint["id"], - } - } - - async with self._cursor(pipeline=True) as cur: - await cur.execute( - """DELETE FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id NOT IN (%s, %s)""", - ( - thread_id, - checkpoint_ns, - checkpoint["id"], - configurable.get("checkpoint_id", ""), - ), - ) - await cur.executemany( - self.UPSERT_CHECKPOINT_BLOBS_SQL, - _dump_blobs( - self.serde, - thread_id, - checkpoint_ns, - copy.pop("channel_values"), # type: ignore[misc] - new_versions, - ), - ) - await cur.execute( - self.UPSERT_CHECKPOINTS_SQL, - ( - thread_id, - checkpoint_ns, - Jsonb(self._dump_checkpoint(copy)), - self._dump_metadata(get_checkpoint_metadata(config, metadata)), - ), - ) - return next_config - - async def aput_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint asynchronously. - - This method saves intermediate writes associated with a checkpoint to the database. - - Args: - config: Configuration of the related checkpoint. - writes: List of writes to store, each as (channel, value) pair. - task_id: Identifier for the task creating the writes. - """ - query = ( - self.UPSERT_CHECKPOINT_WRITES_SQL - if all(w[0] in WRITES_IDX_MAP for w in writes) - else self.INSERT_CHECKPOINT_WRITES_SQL - ) - params = await asyncio.to_thread( - self._dump_writes, - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - task_path, - writes, - ) - async with self._cursor(pipeline=True) as cur: - await cur.executemany(query, params) - - @asynccontextmanager - async def _cursor( - self, *, pipeline: bool = False - ) -> AsyncIterator[AsyncCursor[DictRow]]: - """Create a database cursor as a context manager. - - Args: - pipeline: whether to use pipeline for the DB operations inside the context manager. - Will be applied regardless of whether the AsyncShallowPostgresSaver instance was initialized with a pipeline. - If pipeline mode is not supported, will fall back to using transaction context manager. - """ - async with _ainternal.get_connection(self.conn) as conn: - if self.pipe: - # a connection in pipeline mode can be used concurrently - # in multiple threads/coroutines, but only one cursor can be - # used at a time - try: - async with conn.cursor(binary=True, row_factory=dict_row) as cur: - yield cur - finally: - if pipeline: - await self.pipe.sync() - elif pipeline: - # a connection not in pipeline mode can only be used by one - # thread/coroutine at a time, so we acquire a lock - if self.supports_pipeline: - async with ( - self.lock, - conn.pipeline(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - # Use connection's transaction context manager when pipeline mode not supported - async with ( - self.lock, - conn.transaction(), - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - else: - async with ( - self.lock, - conn.cursor(binary=True, row_factory=dict_row) as cur, - ): - yield cur - - def list( - self, - config: Optional[RunnableConfig], - *, - filter: Optional[dict[str, Any]] = None, - before: Optional[RunnableConfig] = None, - limit: Optional[int] = None, - ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. - - This method retrieves a list of checkpoint tuples from the Postgres database based - on the provided config. For ShallowPostgresSaver, this method returns a list with - ONLY the most recent checkpoint. - """ - aiter_ = self.alist(config, filter=filter, before=before, limit=limit) - while True: - try: - yield asyncio.run_coroutine_threadsafe( - anext(aiter_), # noqa: F821 - self.loop, - ).result() - except StopAsyncIteration: - break - - def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - """Get a checkpoint tuple from the database. - - This method retrieves a checkpoint tuple from the Postgres database based on the - provided config (matching the thread ID in the config). - - Args: - config: The config to use for retrieving the checkpoint. - - Returns: - Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. - """ - try: - # check if we are in the main thread, only bg threads can block - # we don't check in other methods to avoid the overhead - if asyncio.get_running_loop() is self.loop: - raise asyncio.InvalidStateError( - "Synchronous calls to AsyncShallowPostgresSaver are only allowed from a " - "different thread. From the main thread, use the async interface." - "For example, use `await checkpointer.aget_tuple(...)` or `await " - "graph.ainvoke(...)`." - ) - except RuntimeError: - pass - return asyncio.run_coroutine_threadsafe( - self.aget_tuple(config), self.loop - ).result() - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - """Save a checkpoint to the database. - - This method saves a checkpoint to the Postgres database. The checkpoint is associated - with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent - checkpoint and overwrites a previous checkpoint, if it exists. - - Args: - config: The config to associate with the checkpoint. - checkpoint: The checkpoint to save. - metadata: Additional metadata to save with the checkpoint. - new_versions: New channel versions as of this write. - - Returns: - RunnableConfig: Updated configuration after storing the checkpoint. - """ - return asyncio.run_coroutine_threadsafe( - self.aput(config, checkpoint, metadata, new_versions), self.loop - ).result() - - def put_writes( - self, - config: RunnableConfig, - writes: Sequence[tuple[str, Any]], - task_id: str, - task_path: str = "", - ) -> None: - """Store intermediate writes linked to a checkpoint. - - This method saves intermediate writes associated with a checkpoint to the database. - - Args: - config: Configuration of the related checkpoint. - writes: List of writes to store, each as (channel, value) pair. - task_id: Identifier for the task creating the writes. - task_path: Path of the task creating the writes. - """ - return asyncio.run_coroutine_threadsafe( - self.aput_writes(config, writes, task_id, task_path), self.loop - ).result() diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 2beffb3a5..45ce0f082 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -17,10 +17,7 @@ from langgraph.checkpoint.base import ( create_checkpoint, empty_checkpoint, ) -from langgraph.checkpoint.postgres.aio import ( - AsyncPostgresSaver, - AsyncShallowPostgresSaver, -) +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from tests.conftest import DEFAULT_POSTGRES_URI @@ -111,41 +108,11 @@ async def _base_saver(): await conn.execute(f"DROP DATABASE {database}") -@asynccontextmanager -async def _shallow_saver(): - """Fixture for shallow connection mode testing.""" - 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: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI + database, - autocommit=True, - prepare_threshold=0, - row_factory=dict_row, - ) as conn: - checkpointer = AsyncShallowPostgresSaver(conn) - 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 _saver(name: str): if name == "base": async with _base_saver() as saver: yield saver - elif name == "shallow": - async with _shallow_saver() as saver: - yield saver elif name == "pool": async with _pool_saver() as saver: yield saver @@ -205,7 +172,7 @@ def test_data(): } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) async def test_combined_metadata(saver_name: str, test_data) -> None: async with _saver(saver_name) as saver: config = { @@ -232,7 +199,7 @@ async def test_combined_metadata(saver_name: str, test_data) -> None: } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) async def test_asearch(saver_name: str, test_data) -> None: async with _saver(saver_name) as saver: configs = test_data["configs"] @@ -283,7 +250,7 @@ async def test_asearch(saver_name: str, test_data) -> None: } == {"", "inner"} -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) async def test_null_chars(saver_name: str, test_data) -> None: async with _saver(saver_name) as saver: config = await saver.aput( diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index b78f38e9a..4409fa2d6 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -18,7 +18,7 @@ from langgraph.checkpoint.base import ( create_checkpoint, empty_checkpoint, ) -from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver +from langgraph.checkpoint.postgres import PostgresSaver from tests.conftest import DEFAULT_POSTGRES_URI @@ -97,37 +97,11 @@ def _base_saver(): conn.execute(f"DROP DATABASE {database}") -@contextmanager -def _shallow_saver(): - """Fixture for regular connection mode testing with a shallow checkpointer.""" - 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: - with Connection.connect( - DEFAULT_POSTGRES_URI + database, - autocommit=True, - prepare_threshold=0, - row_factory=dict_row, - ) as conn: - checkpointer = ShallowPostgresSaver(conn) - 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 _saver(name: str): if name == "base": with _base_saver() as saver: yield saver - elif name == "shallow": - with _shallow_saver() as saver: - yield saver elif name == "pool": with _pool_saver() as saver: yield saver @@ -187,7 +161,7 @@ def test_data(): } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) def test_combined_metadata(saver_name: str, test_data) -> None: with _saver(saver_name) as saver: config = { @@ -214,7 +188,7 @@ def test_combined_metadata(saver_name: str, test_data) -> None: } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) def test_search(saver_name: str, test_data) -> None: with _saver(saver_name) as saver: configs = test_data["configs"] @@ -263,7 +237,7 @@ def test_search(saver_name: str, test_data) -> None: } == {"", "inner"} -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) def test_null_chars(saver_name: str, test_data) -> None: with _saver(saver_name) as saver: config = saver.put( diff --git a/libs/langgraph/.claude/settings.local.json b/libs/langgraph/.claude/settings.local.json new file mode 100644 index 000000000..67e11f281 --- /dev/null +++ b/libs/langgraph/.claude/settings.local.json @@ -0,0 +1,9 @@ +{ + "permissions": { + "allow": [ + "Bash(rg:*)", + "Bash(python:*)" + ], + "deny": [] + } +} \ No newline at end of file diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 8ca3a8daf..75f83cbdc 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -19,10 +19,8 @@ from tests.conftest_checkpointer import ( _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, @@ -91,12 +89,6 @@ def checkpointer_postgres(): yield checkpointer -@pytest.fixture(scope="function") -def checkpointer_postgres_shallow(): - with _checkpointer_postgres_shallow() as checkpointer: - yield checkpointer - - @pytest.fixture(scope="function") def checkpointer_postgres_pipe(): with _checkpointer_postgres_pipe() as checkpointer: @@ -124,9 +116,6 @@ async def awith_checkpointer( elif checkpointer_name == "postgres_aio": async with _checkpointer_postgres_aio() as checkpointer: yield checkpointer - elif checkpointer_name == "postgres_aio_shallow": - async with _checkpointer_postgres_aio_shallow() as checkpointer: - yield checkpointer elif checkpointer_name == "postgres_aio_pipe": async with _checkpointer_postgres_aio_pipe() as checkpointer: yield checkpointer @@ -275,7 +264,6 @@ ALL_CHECKPOINTERS_SYNC = [ "postgres", "postgres_pipe", "postgres_pool", - "postgres_shallow", ] ALL_CHECKPOINTERS_ASYNC = [ "memory", @@ -283,5 +271,4 @@ ALL_CHECKPOINTERS_ASYNC = [ "postgres_aio", "postgres_aio_pipe", "postgres_aio_pool", - "postgres_aio_shallow", ] diff --git a/libs/langgraph/tests/conftest_checkpointer.py b/libs/langgraph/tests/conftest_checkpointer.py index e43773586..84bd16a1c 100644 --- a/libs/langgraph/tests/conftest_checkpointer.py +++ b/libs/langgraph/tests/conftest_checkpointer.py @@ -6,11 +6,8 @@ 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.postgres import PostgresSaver +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver @@ -58,25 +55,6 @@ def _checkpointer_postgres(): 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]}" @@ -150,31 +128,6 @@ async def _checkpointer_postgres_aio(): 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): @@ -234,12 +187,10 @@ __all__ = [ "_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_large_cases.py b/libs/langgraph/tests/test_large_cases.py index e0b41f48a..5e56d0062 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -110,8 +110,6 @@ def test_invoke_two_processes_in_out_interrupt( snapshot = app.get_state(thread2) assert snapshot.next == () - if "shallow" in checkpointer_name: - return # list history history = [c for c in app.get_state_history(thread1)] @@ -776,10 +774,7 @@ def test_conditional_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -839,10 +834,7 @@ def test_conditional_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -971,10 +963,7 @@ def test_conditional_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1038,10 +1027,7 @@ def test_conditional_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1095,10 +1081,7 @@ def test_conditional_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1227,10 +1210,7 @@ def test_conditional_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1294,10 +1274,7 @@ def test_conditional_graph( }, "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1660,10 +1637,7 @@ def test_conditional_state_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1713,10 +1687,7 @@ def test_conditional_state_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1799,10 +1770,7 @@ def test_conditional_state_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1862,10 +1830,7 @@ def test_conditional_state_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1915,10 +1880,7 @@ def test_conditional_state_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2001,10 +1963,7 @@ def test_conditional_state_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2045,10 +2004,7 @@ def test_conditional_state_graph( "writes": None, "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2096,10 +2052,7 @@ def test_conditional_state_graph( }, "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2168,10 +2121,7 @@ def test_conditional_state_graph( }, "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2242,10 +2192,7 @@ def test_conditional_state_graph( }, "thread_id": "4", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2314,10 +2261,7 @@ def test_conditional_state_graph( }, "thread_id": "4", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -3003,10 +2947,7 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3068,10 +3009,7 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3188,10 +3126,7 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3250,10 +3185,7 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3340,10 +3272,7 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3399,10 +3328,7 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3517,10 +3443,7 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3579,10 +3502,7 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3876,10 +3796,7 @@ def test_message_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -3928,10 +3845,7 @@ def test_message_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4022,10 +3936,7 @@ def test_message_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4074,10 +3985,7 @@ def test_message_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4150,10 +4058,7 @@ def test_message_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4208,10 +4113,7 @@ def test_message_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4302,10 +4204,7 @@ def test_message_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4355,10 +4254,7 @@ def test_message_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4408,10 +4304,7 @@ def test_message_graph( "writes": {"tools": UnsortedSequence("ai", "an extra message")}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4708,10 +4601,7 @@ def test_root_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4760,10 +4650,7 @@ def test_root_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4855,10 +4742,7 @@ def test_root_graph( }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4908,10 +4792,7 @@ def test_root_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4984,10 +4865,7 @@ def test_root_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5042,10 +4920,7 @@ def test_root_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5137,10 +5012,7 @@ def test_root_graph( }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5189,10 +5061,7 @@ def test_root_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5242,10 +5111,7 @@ def test_root_graph( "writes": {"tools": UnsortedSequence("ai", "an extra message")}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5326,10 +5192,7 @@ def test_root_graph( "writes": {"tools": UnsortedSequence("ai", "an extra message")}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(new_app.checkpointer.list(config, limit=2))[-1].config + parent_config=(list(new_app.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5684,23 +5547,22 @@ def test_dynamic_interrupt( ], } - if "shallow" not in checkpointer_name: - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -5734,10 +5596,7 @@ def test_dynamic_interrupt( "writes": None, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=( Interrupt( @@ -5769,10 +5628,7 @@ def test_dynamic_interrupt( "writes": {}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -5875,23 +5731,22 @@ def test_copy_checkpoint( Interrupt(value="Just because...", resumable=True, ns=[AnyStr("tool_two:")]) ], } - if "shallow" not in checkpointer_name: - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️ one", "market": "DE"}, @@ -5931,10 +5786,7 @@ def test_copy_checkpoint( "writes": None, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config + parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config ), interrupts=( Interrupt( @@ -5945,8 +5797,6 @@ def test_copy_checkpoint( ), ) - if "shallow" in checkpointer_name: - return # clear the interrupt and next tasks tool_two.update_state(thread1, None, as_node="__copy__") @@ -6094,28 +5944,27 @@ def test_dynamic_interrupt_subgraph( ], } - if "shallow" not in checkpointer_name: - assert [ - c.metadata - for c in tool_two.checkpointer.list( - {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} - ) - ] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] + assert [ + c.metadata + for c in tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + ) + ] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -6155,10 +6004,7 @@ def test_dynamic_interrupt_subgraph( "writes": None, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list( + parent_config=(list( tool_two.checkpointer.list( {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 ) @@ -6194,10 +6040,7 @@ def test_dynamic_interrupt_subgraph( "writes": {}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list( + parent_config=(list( tool_two.checkpointer.list( {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 ) @@ -6261,25 +6104,24 @@ def test_start_branch_then( "market": "DE", } - if "shallow" not in checkpointer_name: - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "assistant_id": "a", + "thread_id": "1", + }, + ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -6301,10 +6143,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -6333,10 +6172,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -6367,10 +6203,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config ), interrupts=(), ) @@ -6399,10 +6232,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config ), interrupts=(), ) @@ -6433,10 +6263,7 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config ), interrupts=(), ) @@ -6462,10 +6289,7 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config ), interrupts=(), ) @@ -6494,10 +6318,7 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config ), interrupts=(), ) @@ -6874,10 +6695,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -6905,10 +6723,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -6938,10 +6753,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config ), interrupts=(), ) @@ -6969,10 +6781,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config ), interrupts=(), ) @@ -7010,10 +6819,7 @@ def test_branch_then( "writes": {"tool_two_slow": {"my_key": " slow"}}, "thread_id": "11", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -7042,10 +6848,7 @@ def test_branch_then( "writes": {"tool_two_slow": {"my_key": "er"}}, "thread_id": "11", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -7083,10 +6886,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "21", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -7114,10 +6914,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "21", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread1, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config ), interrupts=(), ) @@ -7147,10 +6944,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "22", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config ), interrupts=(), ) @@ -7178,10 +6972,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "22", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread2, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config ), interrupts=(), ) @@ -7237,10 +7028,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "23", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config ), interrupts=(), ) @@ -7268,10 +7056,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "23", }, - parent_config=( - None - if "shallow" in checkpointer_name - else list(tool_two.checkpointer.list(thread3, limit=2))[-1].config + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config ), interrupts=(), ) @@ -7743,10 +7528,7 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7807,10 +7589,7 @@ def test_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -7841,10 +7620,7 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7888,10 +7664,7 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7967,9 +7740,6 @@ def test_nested_graph_state( ), ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] - assert history == expected_history # get_state_history for a subgraph returns its checkpoints @@ -8006,10 +7776,7 @@ def test_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -8111,9 +7878,6 @@ def test_nested_graph_state( ), ] - if "shallow" in checkpointer_name: - expected_child_history = expected_child_history[:1] - assert child_history == expected_child_history # resume @@ -8140,10 +7904,7 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8177,10 +7938,7 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8328,8 +8086,6 @@ def test_nested_graph_state( interrupts=(), ), ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] assert actual_history == expected_history # test looking up parent state by checkpoint ID @@ -8435,10 +8191,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8491,10 +8244,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -8553,10 +8303,7 @@ def test_doubly_nested_graph_state( "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -8639,10 +8386,7 @@ def test_doubly_nested_graph_state( ], }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -8687,10 +8431,7 @@ def test_doubly_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -8721,10 +8462,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8769,10 +8507,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8784,8 +8519,6 @@ def test_doubly_nested_graph_state( ) ) - if "shallow" in checkpointer_name: - return # get outer graph history outer_history = list(app.get_state_history(config)) @@ -9580,10 +9313,7 @@ def test_send_react_interrupt( "thread_id": "2", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -9644,10 +9374,7 @@ def test_send_react_interrupt( "thread_id": "2", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -9742,10 +9469,7 @@ def test_send_react_interrupt( "thread_id": "3", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "3", "checkpoint_ns": "", @@ -9834,10 +9558,7 @@ def test_send_react_interrupt( "thread_id": "3", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "3", "checkpoint_ns": "", @@ -10054,10 +9775,7 @@ def test_send_react_interrupt_control( "thread_id": "2", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -10118,10 +9836,7 @@ def test_send_react_interrupt_control( "thread_id": "2", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -10308,10 +10023,7 @@ def test_weather_subgraph( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -10401,10 +10113,7 @@ def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": "", @@ -10451,10 +10160,7 @@ def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), @@ -10508,10 +10214,7 @@ def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": "", @@ -10565,10 +10268,7 @@ def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 91e4886ef..3a16a1034 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -109,9 +109,6 @@ async def test_invoke_two_processes_in_out_interrupt( snapshot = await app.aget_state(thread2) assert snapshot.next == () - if "shallow" in checkpointer_name: - return - # list history history = [c async for c in app.aget_state_history(thread1)] assert history == [ @@ -852,13 +849,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -911,13 +904,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1045,13 +1034,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1121,13 +1106,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1180,13 +1161,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1314,13 +1291,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1390,13 +1363,9 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1777,13 +1746,9 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -1832,10 +1797,7 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -1920,10 +1882,7 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -1989,10 +1948,7 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -2044,13 +2000,9 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -2132,13 +2084,9 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, interrupts=(), ) @@ -2773,10 +2721,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -2832,10 +2777,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -2947,10 +2889,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3002,10 +2941,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3090,10 +3026,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3149,10 +3082,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3264,10 +3194,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3319,10 +3246,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: }, "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3585,10 +3509,7 @@ async def test_message_graph(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3640,10 +3561,7 @@ async def test_message_graph(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3731,10 +3649,7 @@ async def test_message_graph(checkpointer_name: str) -> None: }, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -3780,10 +3695,7 @@ async def test_message_graph(checkpointer_name: str) -> None: "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [ + parent_config=([ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) ][-1].config ), @@ -4063,25 +3975,24 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "my_key": "value", "market": "DE", } - if "shallow" not in checkpointer_name: - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "assistant_id": "a", + "thread_id": "1", + }, + ] assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value", "market": "DE"}, @@ -4103,13 +4014,9 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "a", "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, interrupts=(), ) # resume, for same result as above @@ -4137,13 +4044,9 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "a", "thread_id": "1", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, interrupts=(), ) @@ -4173,10 +4076,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "a", "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ -1 ].config ), @@ -4207,10 +4107,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "a", "thread_id": "2", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ -1 ].config ), @@ -4243,10 +4140,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "b", "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ -1 ].config ), @@ -4274,10 +4168,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "b", "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ -1 ].config ), @@ -4308,10 +4199,7 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "assistant_id": "b", "thread_id": "3", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ -1 ].config ), @@ -4851,10 +4739,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "11", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -4884,10 +4769,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"finish": {"my_key": " finished"}}, "thread_id": "11", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -4919,10 +4801,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "12", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ -1 ].config ), @@ -4952,10 +4831,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"finish": {"my_key": " finished"}}, "thread_id": "12", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ -1 ].config ), @@ -4995,10 +4871,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "21", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -5028,10 +4901,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"finish": {"my_key": " finished"}}, "thread_id": "21", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -5063,10 +4933,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "22", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ -1 ].config ), @@ -5096,10 +4963,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"finish": {"my_key": " finished"}}, "thread_id": "22", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ -1 ].config ), @@ -5153,7 +5017,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "23", }, - parent_config=(None if "shallow" in checkpointer_name else uconfig), + parent_config=(uconfig), interrupts=(), ) # resume, for same result as above @@ -5180,10 +5044,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "writes": {"finish": {"my_key": " finished"}}, "thread_id": "23", }, - parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ -1 ].config ), @@ -5273,10 +5134,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5337,20 +5195,16 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } - ), + }, interrupts=(), ), ), @@ -5371,10 +5225,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5418,10 +5269,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5497,9 +5345,6 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] - assert history == expected_history # get_state_history for a subgraph returns its checkpoints @@ -5538,10 +5383,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -5643,9 +5485,6 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), ] - if "shallow" in checkpointer_name: - expected_child_history = expected_child_history[:1] - assert child_history == expected_child_history # resume @@ -5672,10 +5511,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5711,10 +5547,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5865,8 +5698,6 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: interrupts=(), ), ] - if "shallow" in checkpointer_name: - expected_history = expected_history[:1] assert actual_history == expected_history # test looking up parent state by checkpoint ID @@ -5971,10 +5802,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -6027,10 +5855,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -6091,10 +5916,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: ], }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -6179,10 +6001,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: ], }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -6231,10 +6050,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -6265,10 +6081,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -6318,10 +6131,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -6333,8 +6143,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: ) ) - if "shallow" in checkpointer_name: - return # get outer graph history outer_history = [c async for c in app.aget_state_history(config)] @@ -7138,10 +6946,7 @@ async def test_weather_subgraph( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7235,10 +7040,7 @@ async def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": "", @@ -7285,10 +7087,7 @@ async def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), @@ -7342,10 +7141,7 @@ async def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": "", @@ -7402,10 +7198,7 @@ async def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0f66c63be..cc104c288 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1099,8 +1099,6 @@ def test_invoke_checkpoint_two( def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") checkpointer: BaseCheckpointSaver = request.getfixturevalue( f"checkpointer_{checkpointer_name}" @@ -1203,9 +1201,6 @@ def test_pending_writes_resume( "value": 6 } - if "shallow" in checkpointer_name: - assert len(list(checkpointer.list(thread1))) == 1 - return # check all final checkpoints checkpoints = [c for c in checkpointer.list(thread1)] @@ -1497,8 +1492,6 @@ def test_send_sequences() -> None: def test_imp_task( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") mapper_calls = 0 @@ -1594,8 +1587,6 @@ def test_imp_task( def test_imp_nested( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -1667,8 +1658,6 @@ def test_imp_nested( def test_imp_stream_order( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @@ -1789,8 +1778,6 @@ def test_invoke_checkpoint_three( assert state.values.get("total") == 5 assert state.next == () - if "shallow" in checkpointer_name: - return assert len(list(app.get_state_history(thread_1, limit=1))) == 1 # list all checkpoints for thread 1 @@ -2398,10 +2385,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt.update_state(config, {"docs": ["doc5"]}) - expected_parent_config = ( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + expected_parent_config = (list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ) assert app_w_interrupt.get_state(config) == StateSnapshot( values={ @@ -2677,10 +2661,7 @@ def test_in_one_fan_out_state_graph_defer_node( ] app_w_interrupt.update_state(config, {"docs": ["doc5"]}) - expected_parent_config = ( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + expected_parent_config = (list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ) assert app_w_interrupt.get_state(config) == StateSnapshot( values={ @@ -2955,10 +2936,7 @@ def test_in_one_fan_out_state_graph_then_defer_node( ] app_w_interrupt.update_state(config, {"docs": ["doc5"]}) - expected_parent_config = ( - None - if "shallow" in checkpointer_name - else list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + expected_parent_config = (list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ) assert app_w_interrupt.get_state(config) == StateSnapshot( values={ @@ -3890,8 +3868,6 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: def test_subgraph_checkpoint_true( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Unsupported combo") checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) @@ -3958,8 +3934,6 @@ def test_subgraph_checkpoint_true( def test_subgraph_checkpoint_true_interrupt( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Unsupported combo") checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) @@ -4139,8 +4113,6 @@ def test_stream_buffering_single_node( def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Unsupported combo") checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) @@ -4306,8 +4278,6 @@ def test_nested_graph_interrupts_parallel( def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Unsupported combo") checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) @@ -5498,10 +5468,7 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) "parents": {}, }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { + parent_config=({ "configurable": { "thread_id": "1", "checkpoint_ns": "", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4cbdc465e..da2ab68dd 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -595,23 +595,22 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: ) }, ] - if "shallow" not in checkpointer_name: - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] tup = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -640,9 +639,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -673,9 +670,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -793,25 +788,22 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: ) }, ] - if "shallow" not in checkpointer_name: - assert [ - c.metadata async for c in tool_two.checkpointer.alist(thread1root) - ] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] tup = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, @@ -846,11 +838,9 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in tool_two.checkpointer.alist(thread1root, limit=2) - ][-1].config + [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ + -1 + ].config ), interrupts=( Interrupt( @@ -879,11 +869,9 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [ - c async for c in tool_two.checkpointer.alist(thread1root, limit=2) - ][-1].config + [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ + -1 + ].config ), interrupts=(), ) @@ -998,23 +986,22 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: ], } - if "shallow" not in checkpointer_name: - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] tup = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( @@ -1053,9 +1040,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 ].config ), @@ -1068,10 +1053,6 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: ), ) - if "shallow" in checkpointer_name: - # shallow checkpointer doesn't support copy - return - # clear the interrupt and next tasks await tool_two.aupdate_state(thread1, None, as_node="__copy__") # interrupt is cleared, next task is kept @@ -1998,9 +1979,6 @@ async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) async def test_pending_writes_resume( checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") - class State(TypedDict): value: Annotated[int, operator.add] @@ -2106,10 +2084,6 @@ async def test_pending_writes_resume( None, thread1, checkpoint_during=checkpoint_during ) == {"value": 6} - if "shallow" in checkpointer_name: - assert len([c async for c in checkpointer.alist(thread1)]) == 1 - return - # check all final checkpoints checkpoints = [c async for c in checkpointer.alist(thread1)] # we should have 3 @@ -2499,9 +2473,6 @@ async def test_send_sequences(checkpointer_name: str) -> None: @NEEDS_CONTEXTVARS @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: - pytest.skip("Checkpointing during execution not supported") - async with awith_checkpointer(checkpointer_name) as checkpointer: mapper_calls = 0 @@ -2562,9 +2533,6 @@ async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None @NEEDS_CONTEXTVARS @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: - pytest.skip("Checkpointing during execution not supported") - async def mynode(input: list[str]) -> list[str]: return [it + "a" for it in input] @@ -2637,9 +2605,6 @@ async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> No @NEEDS_CONTEXTVARS @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: - pytest.skip("Checkpointing during execution not supported") - async with awith_checkpointer(checkpointer_name) as checkpointer: mapper_calls = 0 mapper_cancels = 0 @@ -2700,9 +2665,6 @@ async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) async def test_imp_sync_from_async( checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") - async with awith_checkpointer(checkpointer_name) as checkpointer: @task() @@ -2743,9 +2705,6 @@ async def test_imp_sync_from_async( async def test_imp_stream_order( checkpointer_name: str, checkpoint_during: bool ) -> None: - if not checkpoint_during and "shallow" in checkpointer_name: - pytest.skip("Checkpointing during execution not supported") - async with awith_checkpointer(checkpointer_name) as checkpointer: @task() @@ -3324,9 +3283,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -3388,9 +3345,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -3484,9 +3439,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "3", "checkpoint_ns": "", @@ -3576,9 +3529,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "3", "checkpoint_ns": "", @@ -3794,9 +3745,7 @@ async def test_send_react_interrupt_control( }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -3858,9 +3807,7 @@ async def test_send_react_interrupt_control( }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -4098,9 +4045,6 @@ async def test_invoke_checkpoint_three( assert state.values.get("total") == 5 assert state.next == () - if "shallow" in checkpointer_name: - return - assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 # list all checkpoints for thread 1 thread_1_history = [c async for c in app.aget_state_history(thread_1)] @@ -5019,9 +4963,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -6738,9 +6680,7 @@ async def test_parent_command(checkpointer_name: str) -> None: }, created_at=AnyStr(), parent_config=( - None - if "shallow" in checkpointer_name - else { + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7285,9 +7225,6 @@ async def test_checkpoint_recovery_async(checkpointer_name: str): result = await graph.ainvoke({"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 = [c async for c in graph.aget_state_history(config)] assert len(history) == 6 # Initial + failed attempt + successful attempt