diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 8e4595a37..691654e6b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -23,6 +23,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 @@ -456,4 +457,4 @@ class PostgresSaver(BasePostgresSaver): ) -__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"] +__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 3d6396f29..9fc7673b2 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -23,6 +23,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 @@ -559,4 +560,4 @@ class AsyncPostgresSaver(BasePostgresSaver): ).result() -__all__ = ["AsyncPostgresSaver", "Conn"] +__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py new file mode 100644 index 000000000..af2aacef8 --- /dev/null +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py @@ -0,0 +1,959 @@ +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: Checkpoint = { + **value["checkpoint"], + "channel_values": self._load_blobs(value["channel_values"]), + "pending_sends": [ + self.serde.loads_typed((t.decode(), v)) + for t, v in value["pending_sends"] + ] + if value["pending_sends"] + else [], + } + yield CheckpointTuple( + config={ + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": checkpoint["id"], + } + }, + checkpoint=checkpoint, + 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: Checkpoint = { + **value["checkpoint"], + "channel_values": self._load_blobs(value["channel_values"]), + "pending_sends": [ + self.serde.loads_typed((t.decode(), v)) + for t, v in value["pending_sends"] + ] + if value["pending_sends"] + else [], + } + return CheckpointTuple( + config={ + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint["id"], + } + }, + checkpoint=checkpoint, + 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(copy), + Jsonb(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: Checkpoint = { + **value["checkpoint"], + "channel_values": self._load_blobs(value["channel_values"]), + "pending_sends": [ + self.serde.loads_typed((t.decode(), v)) + for t, v in value["pending_sends"] + ] + if value["pending_sends"] + else [], + } + yield CheckpointTuple( + config={ + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": checkpoint["id"], + } + }, + checkpoint=checkpoint, + 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: Checkpoint = { + **value["checkpoint"], + "channel_values": self._load_blobs(value["channel_values"]), + "pending_sends": [ + self.serde.loads_typed((t.decode(), v)) + for t, v in value["pending_sends"] + ] + if value["pending_sends"] + else [], + } + return CheckpointTuple( + config={ + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint["id"], + } + }, + checkpoint=checkpoint, + 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(copy), + Jsonb(get_checkpoint_metadata(config, metadata)), + ), + ) + return next_config + + async def aput_writes( + self, + config: RunnableConfig, + writes: Sequence[tuple[str, Any]], + task_id: str, + task_path: str = "", + ) -> None: + """Store intermediate writes linked to a checkpoint asynchronously. + + This method saves intermediate writes associated with a checkpoint to the database. + + Args: + config: Configuration of the related checkpoint. + writes: List of writes to store, each as (channel, value) pair. + task_id: Identifier for the task creating the writes. + """ + query = ( + self.UPSERT_CHECKPOINT_WRITES_SQL + if all(w[0] in WRITES_IDX_MAP for w in writes) + else self.INSERT_CHECKPOINT_WRITES_SQL + ) + params = await asyncio.to_thread( + self._dump_writes, + config["configurable"]["thread_id"], + config["configurable"]["checkpoint_ns"], + config["configurable"]["checkpoint_id"], + task_id, + task_path, + writes, + ) + async with self._cursor(pipeline=True) as cur: + await cur.executemany(query, params) + + @asynccontextmanager + async def _cursor( + self, *, pipeline: bool = False + ) -> AsyncIterator[AsyncCursor[DictRow]]: + """Create a database cursor as a context manager. + + Args: + pipeline: whether to use pipeline for the DB operations inside the context manager. + Will be applied regardless of whether the AsyncShallowPostgresSaver instance was initialized with a pipeline. + If pipeline mode is not supported, will fall back to using transaction context manager. + """ + async with _ainternal.get_connection(self.conn) as conn: + if self.pipe: + # a connection in pipeline mode can be used concurrently + # in multiple threads/coroutines, but only one cursor can be + # used at a time + try: + async with conn.cursor(binary=True, row_factory=dict_row) as cur: + yield cur + finally: + if pipeline: + await self.pipe.sync() + elif pipeline: + # a connection not in pipeline mode can only be used by one + # thread/coroutine at a time, so we acquire a lock + if self.supports_pipeline: + async with ( + self.lock, + conn.pipeline(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): + yield cur + else: + # Use connection's transaction context manager when pipeline mode not supported + async with ( + self.lock, + conn.transaction(), + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): + yield cur + else: + async with ( + self.lock, + conn.cursor(binary=True, row_factory=dict_row) as cur, + ): + yield cur + + def list( + self, + config: Optional[RunnableConfig], + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[CheckpointTuple]: + """List checkpoints from the database. + + This method retrieves a list of checkpoint tuples from the Postgres database based + on the provided config. For ShallowPostgresSaver, this method returns a list with + ONLY the most recent checkpoint. + """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), # noqa: F821 + self.loop, + ).result() + except StopAsyncIteration: + break + + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + """Get a checkpoint tuple from the database. + + This method retrieves a checkpoint tuple from the Postgres database based on the + provided config (matching the thread ID in the config). + + Args: + config: The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + try: + # check if we are in the main thread, only bg threads can block + # we don't check in other methods to avoid the overhead + if asyncio.get_running_loop() is self.loop: + raise asyncio.InvalidStateError( + "Synchronous calls to AsyncShallowPostgresSaver are only allowed from a " + "different thread. From the main thread, use the async interface." + "For example, use `await checkpointer.aget_tuple(...)` or `await " + "graph.ainvoke(...)`." + ) + except RuntimeError: + pass + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() + + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + """Save a checkpoint to the database. + + This method saves a checkpoint to the Postgres database. The checkpoint is associated + with the provided config. For AsyncShallowPostgresSaver, this method saves ONLY the most recent + checkpoint and overwrites a previous checkpoint, if it exists. + + Args: + config: The config to associate with the checkpoint. + checkpoint: The checkpoint to save. + metadata: Additional metadata to save with the checkpoint. + new_versions: New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, + config: RunnableConfig, + writes: Sequence[tuple[str, Any]], + task_id: str, + task_path: str = "", + ) -> None: + """Store intermediate writes linked to a checkpoint. + + This method saves intermediate writes associated with a checkpoint to the database. + + Args: + config: Configuration of the related checkpoint. + writes: List of writes to store, each as (channel, value) pair. + task_id: Identifier for the task creating the writes. + task_path: Path of the task creating the writes. + """ + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id, task_path), self.loop + ).result() diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 59669d17c..f0196f845 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -15,7 +15,10 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, ) -from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.checkpoint.postgres.aio import ( + AsyncPostgresSaver, + AsyncShallowPostgresSaver, +) from langgraph.checkpoint.serde.types import TASKS from tests.checkpoint_utils import create_checkpoint, empty_checkpoint from tests.conftest import DEFAULT_POSTGRES_URI @@ -108,11 +111,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 @@ -172,7 +205,7 @@ def test_data(): } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) async def test_combined_metadata(saver_name: str, test_data) -> None: async with _saver(saver_name) as saver: config = { @@ -199,7 +232,7 @@ async def test_combined_metadata(saver_name: str, test_data) -> None: } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) async def test_asearch(saver_name: str, test_data) -> None: async with _saver(saver_name) as saver: configs = test_data["configs"] @@ -250,7 +283,7 @@ async def test_asearch(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(saver_name: str, test_data) -> None: async with _saver(saver_name) as saver: config = await saver.aput( diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index b6eea12bf..3ce48c8da 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -16,7 +16,7 @@ from langgraph.checkpoint.base import ( Checkpoint, CheckpointMetadata, ) -from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver from langgraph.checkpoint.serde.types import TASKS from tests.checkpoint_utils import create_checkpoint, empty_checkpoint from tests.conftest import DEFAULT_POSTGRES_URI @@ -97,11 +97,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 @@ -161,7 +187,7 @@ def test_data(): } -@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"]) +@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"]) def test_combined_metadata(saver_name: str, test_data) -> None: with _saver(saver_name) as saver: config = { @@ -188,7 +214,7 @@ def test_combined_metadata(saver_name: str, test_data) -> None: } -@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"] @@ -237,7 +263,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(