WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 10:17:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove postgres shallow checkpointer (#4813)

Browse Source
This commit is contained in:
Nuno Campos
2025-05-28 14:02:58 -07:00
committed by GitHub
parent dc95f95181 8c11c1155a
commit e384ccacaf
12 changed files with 366 additions and 2024 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+1 -2
View File
@@ -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"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+1 -2
View File
@@ -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"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py
-941
View File
@@ -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()
libs/checkpoint-postgres/tests/test_async.py
+4 -37
View File
@@ -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(
libs/checkpoint-postgres/tests/test_sync.py
+4 -30
View File
@@ -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(
libs/langgraph/.claude/settings.local.json
+9
View File
@@ -0,0 +1,9 @@
{
"permissions": {
"allow": [
"Bash(rg:*)",
"Bash(python:*)"
],
"deny": []
}
}
libs/langgraph/tests/conftest.py
-13
View File
@@ -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",
]
libs/langgraph/tests/conftest_checkpointer.py
+2 -51
View File
@@ -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",
]
libs/langgraph/tests/test_large_cases.py
+165 -465
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_large_cases_async.py
+111 -318
View File
@@ -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:"),
libs/langgraph/tests/test_pregel.py
+4 -37
View File
@@ -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": "",
libs/langgraph/tests/test_pregel_async.py
+65 -128
View File
@@ -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