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

checkpoint postgres: add a shallow checkpointer (#2826)

Browse Source 
This PR adds a "shallow" version of `PostgresSaver` checkpointer that
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.
This commit is contained in:
Vadym Barda
2024-12-20 17:51:20 +00:00
committed by GitHub
parent 1de61f6fce
commit 44ee0199fd
14 changed files with 4322 additions and 1082 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+2 -1
View File
@@ -19,6 +19,7 @@ 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
@@ -396,4 +397,4 @@ class PostgresSaver(BasePostgresSaver):
yield cur
__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"]
__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+2 -1
View File
@@ -19,6 +19,7 @@ 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
@@ -464,4 +465,4 @@ class AsyncPostgresSaver(BasePostgresSaver):
).result()
__all__ = ["AsyncPostgresSaver", "Conn"]
__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py
+918
View File
@@ -0,0 +1,918 @@
import asyncio
import threading
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,
)
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);
""",
]
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_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, idx, channel, type, blob)
VALUES (%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, idx, channel, type, blob)
VALUES (%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, 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:
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 (str): The Postgres connection info string.
pipeline (bool): 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 (RunnableConfig): 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 (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
new_versions (ChannelVersions): 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(metadata),
),
)
return next_config
def put_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint.
This method saves intermediate writes associated with a checkpoint to the Postgres database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (List[Tuple[str, Any]]): List of writes to store.
task_id (str): 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,
writes,
),
)
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
"""Create a database cursor as a context manager.
Args:
pipeline (bool): 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:
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 (str): The Postgres connection info string.
pipeline (bool): 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 (RunnableConfig): 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 (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
new_versions (ChannelVersions): 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(metadata),
),
)
return next_config
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
This method saves intermediate writes associated with a checkpoint to the database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): 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,
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 (bool): 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 (RunnableConfig): 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 (RunnableConfig): The config to associate with the checkpoint.
checkpoint (Checkpoint): The checkpoint to save.
metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
new_versions (ChannelVersions): 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,
) -> None:
"""Store intermediate writes linked to a checkpoint.
This method saves intermediate writes associated with a checkpoint to the database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
task_id (str): Identifier for the task creating the writes.
"""
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id), self.loop
).result()
libs/checkpoint-postgres/tests/test_async.py
+36 -3
View File
@@ -16,7 +16,10 @@ from langgraph.checkpoint.base import (
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
AsyncShallowPostgresSaver,
)
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -103,11 +106,41 @@ 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
@@ -167,7 +200,7 @@ def test_data():
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
async def test_asearch(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
configs = test_data["configs"]
@@ -212,7 +245,7 @@ async def test_asearch(request, saver_name: str, test_data) -> None:
} == {"", "inner"}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
async def test_null_chars(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
config = await saver.aput(
libs/checkpoint-postgres/tests/test_sync.py
+29 -3
View File
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -91,11 +91,37 @@ 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
@@ -155,7 +181,7 @@ def test_data():
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
def test_search(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
configs = test_data["configs"]
@@ -198,7 +224,7 @@ def test_search(saver_name: str, test_data) -> None:
} == {"", "inner"}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
def test_null_chars(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
config = saver.put(
libs/langgraph/tests/__snapshots__/test_large_cases.ambr
+530
View File
File diff suppressed because one or more lines are too long
libs/langgraph/tests/__snapshots__/test_large_cases_async.ambr
+25
View File
@@ -99,6 +99,31 @@
'''
# ---
# name: test_weather_subgraph[postgres_aio_shallow]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
router_node(router_node)
normal_llm_node(normal_llm_node)
weather_graph_model_node(model_node)
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
__end__([<p>__end__</p>]):::last
__start__ --> router_node;
normal_llm_node --> __end__;
weather_graph_weather_node --> __end__;
router_node -.-> normal_llm_node;
router_node -.-> weather_graph_model_node;
router_node -.-> __end__;
subgraph weather_graph
weather_graph_model_node --> weather_graph_weather_node;
end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_weather_subgraph[sqlite_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+166
View File
@@ -2878,6 +2878,19 @@
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_shallow]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query --> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge[sqlite]
'''
graph TD;
@@ -3311,6 +3324,76 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow].1
dict({
'definitions': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/definitions/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite]
'''
graph TD;
@@ -3788,6 +3871,76 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
]),
'title': 'Input',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow].2
dict({
'properties': dict({
'answer': dict({
'title': 'Answer',
'type': 'string',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
}),
'required': list([
'answer',
'docs',
]),
'title': 'Output',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite]
'''
graph TD;
@@ -3923,6 +4076,19 @@
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_shallow]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite]
'''
graph TD;
libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
+136
View File
@@ -934,6 +934,127 @@
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow]
'''
graph TD;
__start__ --> rewrite_query;
analyzer_one --> retriever_one;
qa --> __end__;
retriever_one --> qa;
retriever_two --> qa;
rewrite_query --> analyzer_one;
rewrite_query -.-> retriever_two;
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].1
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].2
dict({
'$defs': dict({
'InnerObject': dict({
'properties': dict({
'yo': dict({
'title': 'Yo',
'type': 'integer',
}),
}),
'required': list([
'yo',
]),
'title': 'InnerObject',
'type': 'object',
}),
}),
'properties': dict({
'answer': dict({
'anyOf': list([
dict({
'type': 'string',
}),
dict({
'type': 'null',
}),
]),
'default': None,
'title': 'Answer',
}),
'docs': dict({
'items': dict({
'type': 'string',
}),
'title': 'Docs',
'type': 'array',
}),
'inner': dict({
'$ref': '#/$defs/InnerObject',
}),
'query': dict({
'title': 'Query',
'type': 'string',
}),
}),
'required': list([
'query',
'inner',
'docs',
]),
'title': 'State',
'type': 'object',
})
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio]
'''
graph TD;
@@ -1362,6 +1483,21 @@
'''
# ---
# name: test_send_react_interrupt_control[postgres_aio_shallow]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([<p>__start__</p>]):::first
agent(agent)
foo([foo]):::last
__start__ --> agent;
agent -.-> foo;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_send_react_interrupt_control[sqlite_aio]
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
libs/langgraph/tests/conftest.py
+64 -4
View File
@@ -13,8 +13,11 @@ from pytest_mock import MockerFixture
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.duckdb import DuckDBSaver
from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
AsyncShallowPostgresSaver,
)
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.store.base import BaseStore
@@ -100,6 +103,25 @@ def checkpointer_postgres():
conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def checkpointer_postgres_shallow():
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ShallowPostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + database
) as checkpointer:
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def checkpointer_postgres_pipe():
database = f"test_{uuid4().hex[:16]}"
@@ -167,6 +189,31 @@ 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):
@@ -240,6 +287,9 @@ 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
@@ -417,20 +467,30 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
raise NotImplementedError(f"Unknown store {store_name}")
ALL_CHECKPOINTERS_SYNC = [
SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"]
REGULAR_CHECKPOINTERS_SYNC = [
"memory",
"sqlite",
"postgres",
"postgres_pipe",
"postgres_pool",
]
ALL_CHECKPOINTERS_ASYNC = [
ALL_CHECKPOINTERS_SYNC = [
*REGULAR_CHECKPOINTERS_SYNC,
*SHALLOW_CHECKPOINTERS_SYNC,
]
SHALLOW_CHECKPOINTERS_ASYNC = ["postgres_aio_shallow"]
REGULAR_CHECKPOINTERS_ASYNC = [
"memory",
"sqlite_aio",
"postgres_aio",
"postgres_aio_pipe",
"postgres_aio_pool",
]
ALL_CHECKPOINTERS_ASYNC = [
*REGULAR_CHECKPOINTERS_ASYNC,
*SHALLOW_CHECKPOINTERS_ASYNC,
]
ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
*ALL_CHECKPOINTERS_ASYNC,
None,
libs/langgraph/tests/test_large_cases.py
+1365 -516
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_large_cases_async.py
+822 -428
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel.py
+37 -10
View File
@@ -78,6 +78,7 @@ from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence
from tests.conftest import (
ALL_CHECKPOINTERS_SYNC,
ALL_STORES_SYNC,
REGULAR_CHECKPOINTERS_SYNC,
SHOULD_CHECK_SNAPSHOTS,
)
from tests.memory_assert import MemorySaverAssertCheckpointMetadata
@@ -624,7 +625,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert step == 2
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
def test_run_from_checkpoint_id_retains_previous_writes(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
) -> None:
@@ -1157,6 +1158,10 @@ def test_pending_writes_resume(
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
assert graph.invoke(None, thread1) == {"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)]
# we should have 3
@@ -1618,6 +1623,9 @@ 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
thread_1_history = [c for c in app.get_state_history(thread_1)]
@@ -2270,6 +2278,11 @@ 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
)
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"query": "analyzed: query: what is weather in sf",
@@ -2277,8 +2290,14 @@ def test_in_one_fan_out_state_graph_waiting_edge(
},
tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),),
next=("qa",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
@@ -2286,7 +2305,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(
"writes": {"retriever_one": {"docs": ["doc5"]}},
"thread_id": "2",
},
parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config,
parent_config=expected_parent_config,
)
assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [
@@ -4670,13 +4689,17 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str)
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(),
)
@@ -5203,11 +5226,15 @@ def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name:
assert state is not None
assert state.values == {"steps": ["start"], "attempt": 1} # input state saved
assert state.next == ("node1",) # Should retry failed node
assert "RuntimeError('Simulated failure')" in state.tasks[0].error
# Retry with updated attempt count
result = graph.invoke({"steps": [], "attempt": 2}, config)
assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}
if "shallow" in checkpointer_name:
return
# Verify checkpoint history shows both attempts
history = list(graph.get_state_history(config))
assert len(history) == 6 # Initial + failed attempt + successful attempt
libs/langgraph/tests/test_pregel_async.py
+190 -116
View File
@@ -75,6 +75,7 @@ from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
ALL_CHECKPOINTERS_ASYNC_PLUS_NONE,
ALL_STORES_ASYNC,
REGULAR_CHECKPOINTERS_ASYNC,
SHOULD_CHECK_SNAPSHOTS,
awith_checkpointer,
awith_store,
@@ -347,22 +348,23 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
)
},
]
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",
},
]
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",
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
@@ -390,9 +392,13 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
"writes": None,
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
parent_config=(
None
if "shallow" in checkpointer_name
else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].config
),
)
# clear the interrupt and next tasks
@@ -412,9 +418,13 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
"writes": {},
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
parent_config=(
None
if "shallow" in checkpointer_name
else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].config
),
)
@@ -524,22 +534,25 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
)
},
]
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",
},
]
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",
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
@@ -573,9 +586,13 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
"writes": None,
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
][-1].config,
parent_config=(
None
if "shallow" in checkpointer_name
else [
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
][-1].config
),
)
# clear the interrupt and next tasks
@@ -595,9 +612,13 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
"writes": {},
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
][-1].config,
parent_config=(
None
if "shallow" in checkpointer_name
else [
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
][-1].config
),
)
@@ -699,22 +720,25 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"my_key": "value ⛰️ one",
"market": "DE",
}
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
{
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"thread_id": "1",
},
{
"parents": {},
"source": "input",
"step": -1,
"writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
"thread_id": "1",
},
]
if "shallow" not in checkpointer_name:
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
{
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"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 ⛰️ one", "market": "DE"},
@@ -742,9 +766,13 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"writes": {"tool_one": {"my_key": " one"}},
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
parent_config=(
None
if "shallow" in checkpointer_name
else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].config
),
)
# clear the interrupt and next tasks
await tool_two.aupdate_state(thread1, None)
@@ -770,9 +798,13 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"writes": {},
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
parent_config=(
None
if "shallow" in checkpointer_name
else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].config
),
)
@@ -1754,6 +1786,10 @@ async def test_pending_writes_resume(
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
assert await graph.ainvoke(None, thread1) == {"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
@@ -1911,7 +1947,7 @@ async def test_pending_writes_resume(
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_run_from_checkpoint_id_retains_previous_writes(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
) -> None:
@@ -2339,7 +2375,7 @@ async def test_imp_stream_order(checkpointer_name: str) -> None:
]
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
@@ -2792,13 +2828,17 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
"thread_id": "2",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(
PregelTask(
id=AnyStr(),
@@ -2875,13 +2915,17 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
"thread_id": "2",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(),
)
@@ -2951,13 +2995,17 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
"thread_id": "3",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(
PregelTask(
id=AnyStr(),
@@ -3060,13 +3108,17 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
"thread_id": "3",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(
PregelTask(
id=AnyStr(),
@@ -3260,13 +3312,17 @@ async def test_send_react_interrupt_control(
"thread_id": "2",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(
PregelTask(
id=AnyStr(),
@@ -3343,13 +3399,17 @@ async def test_send_react_interrupt_control(
"thread_id": "2",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(),
)
@@ -3572,6 +3632,9 @@ 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)]
@@ -4279,13 +4342,17 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
"thread_id": "1",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
)
async with assert_ctx_once():
@@ -5758,13 +5825,17 @@ async def test_parent_command(checkpointer_name: str) -> None:
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
},
),
tasks=(),
)
@@ -6281,6 +6352,9 @@ 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