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

Remove features and dependencies

Browse Source 
- rm langchain_core dependency
- replace callbacks w run tree
- rm Runnable dependency
- rm non-state Graph
- rm managed values
- rm entrypoint/task/call
- rm async methods
- rm shallow checkpointer
- rm messages stream mode
- rm debug flag
- rm remote graph
This commit is contained in:
Nuno Campos
2025-03-01 13:53:27 -08:00
parent 9284b57ba0
commit eb57c06896
67 changed files with 1301 additions and 33581 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+14 -15
View File
@@ -3,7 +3,6 @@ from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
@@ -13,6 +12,7 @@ from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointConfig,
CheckpointMetadata,
CheckpointTuple,
get_checkpoint_id,
@@ -20,7 +20,6 @@ from langgraph.checkpoint.base import (
)
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import BasePostgresSaver
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
from langgraph.checkpoint.serde.base import SerializerProtocol
Conn = _internal.Conn # For backward compatibility
@@ -97,10 +96,10 @@ class PostgresSaver(BasePostgresSaver):
def list(
self,
config: Optional[RunnableConfig],
config: Optional[CheckpointConfig],
*,
filter: Optional[dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
before: Optional[CheckpointConfig] = None,
limit: Optional[int] = None,
) -> Iterator[CheckpointTuple]:
"""List checkpoints from the database.
@@ -109,9 +108,9 @@ class PostgresSaver(BasePostgresSaver):
on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (RunnableConfig): The config to use for listing the checkpoints.
config (CheckpointConfig): The config to use for listing the checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
before (Optional[CheckpointConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
Yields:
@@ -171,7 +170,7 @@ class PostgresSaver(BasePostgresSaver):
self._load_writes(value["pending_writes"]),
)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
def get_tuple(self, config: CheckpointConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
This method retrieves a checkpoint tuple from the Postgres database based on the
@@ -180,7 +179,7 @@ class PostgresSaver(BasePostgresSaver):
for the given thread ID is retrieved.
Args:
config (RunnableConfig): The config to use for retrieving the checkpoint.
config (CheckpointConfig): The config to use for retrieving the checkpoint.
Returns:
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
@@ -254,24 +253,24 @@ class PostgresSaver(BasePostgresSaver):
def put(
self,
config: RunnableConfig,
config: CheckpointConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
) -> CheckpointConfig:
"""Save a checkpoint to the database.
This method saves a checkpoint to the Postgres database. The checkpoint is associated
with the provided config and its parent config (if any).
Args:
config (RunnableConfig): The config to associate with the checkpoint.
config (CheckpointConfig): 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.
CheckpointConfig: Updated configuration after storing the checkpoint.
Examples:
@@ -325,7 +324,7 @@ class PostgresSaver(BasePostgresSaver):
def put_writes(
self,
config: RunnableConfig,
config: CheckpointConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
@@ -335,7 +334,7 @@ class PostgresSaver(BasePostgresSaver):
This method saves intermediate writes associated with a checkpoint to the Postgres database.
Args:
config (RunnableConfig): Configuration of the related checkpoint.
config (CheckpointConfig): 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.
"""
@@ -400,4 +399,4 @@ class PostgresSaver(BasePostgresSaver):
yield cur
__all__ = ["PostgresSaver", "BasePostgresSaver", "ShallowPostgresSaver", "Conn"]
__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py
-24
View File
@@ -1,24 +0,0 @@
"""Shared async utility functions for the Postgres checkpoint & storage classes."""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Union
from psycopg import AsyncConnection
from psycopg.rows import DictRow
from psycopg_pool import AsyncConnectionPool
Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]]
@asynccontextmanager
async def get_connection(
conn: Conn,
) -> AsyncIterator[AsyncConnection[DictRow]]:
if isinstance(conn, AsyncConnection):
yield conn
elif isinstance(conn, AsyncConnectionPool):
async with conn.connection() as conn:
yield conn
else:
raise TypeError(f"Invalid connection type: {type(conn)}")
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
-485
View File
@@ -1,485 +0,0 @@
import asyncio
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
ChannelVersions,
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
get_checkpoint_id,
get_checkpoint_metadata,
)
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
class AsyncPostgresSaver(BasePostgresSaver):
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["AsyncPostgresSaver"]:
"""Create a new AsyncPostgresSaver instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
pipeline (bool): whether to use AsyncPipeline
Returns:
AsyncPostgresSaver: A new AsyncPostgresSaver 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. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): Maximum number of checkpoints to return.
Yields:
AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
"""
where, args = self._search_where(config, filter, before)
query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
async with self._cursor() as cur:
await cur.execute(query, args, binary=True)
async for value in cur:
yield CheckpointTuple(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None
),
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. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
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_id = get_checkpoint_id(config)
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
if checkpoint_id:
args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
else:
args = (thread_id, checkpoint_ns)
where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
async with self._cursor() as cur:
await cur.execute(
self.SELECT_SQL + where,
args,
binary=True,
)
async for value in cur:
return CheckpointTuple(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["checkpoint_id"],
}
},
await asyncio.to_thread(
self._load_checkpoint,
value["checkpoint"],
value["channel_values"],
value["pending_sends"],
),
self._load_metadata(value["metadata"]),
(
{
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None
),
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 and its parent config (if any).
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")
checkpoint_id = configurable.pop(
"checkpoint_id", configurable.pop("thread_ts", None)
)
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.executemany(
self.UPSERT_CHECKPOINT_BLOBS_SQL,
await asyncio.to_thread(
self._dump_blobs,
thread_id,
checkpoint_ns,
copy.pop("channel_values"), # type: ignore[misc]
new_versions,
),
)
await cur.execute(
self.UPSERT_CHECKPOINTS_SQL,
(
thread_id,
checkpoint_ns,
checkpoint["id"],
checkpoint_id,
Jsonb(self._dump_checkpoint(copy)),
self._dump_metadata(get_checkpoint_metadata(config, metadata)),
),
)
return next_config
async def aput_writes(
self,
config: RunnableConfig,
writes: Sequence[tuple[str, Any]],
task_id: str,
task_path: str = "",
) -> None:
"""Store intermediate writes linked to a checkpoint asynchronously.
This method saves intermediate writes associated with a checkpoint to the database.
Args:
config (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,
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 (bool): whether to use pipeline for the DB operations inside the context manager.
Will be applied regardless of whether the AsyncPostgresSaver 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. The checkpoints are ordered by checkpoint ID in descending order (newest first).
Args:
config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
limit (Optional[int]): Maximum number of checkpoints to return.
Yields:
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
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 AsyncPostgresSaver are only allowed from a "
"different thread. From the main thread, use the async interface. "
"For example, use `checkpointer.alist(...)` or `await "
"graph.ainvoke(...)`."
)
except RuntimeError:
pass
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. If the config contains a "checkpoint_id" key, the checkpoint with
the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
for the given thread ID is retrieved.
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 AsyncPostgresSaver 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 and its parent config (if any).
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,
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 (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.
task_path (str): Path of the task creating the writes.
"""
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id, task_path), self.loop
).result()
__all__ = ["AsyncPostgresSaver", "AsyncShallowPostgresSaver", "Conn"]
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+3 -3
View File
@@ -2,7 +2,6 @@ import random
from collections.abc import Sequence
from typing import Any, Optional, cast
from langchain_core.runnables import RunnableConfig
from psycopg.types.json import Jsonb
from langgraph.checkpoint.base import (
@@ -10,6 +9,7 @@ from langgraph.checkpoint.base import (
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
CheckpointConfig,
CheckpointMetadata,
get_checkpoint_id,
)
@@ -262,9 +262,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
def _search_where(
self,
config: Optional[RunnableConfig],
config: Optional[CheckpointConfig],
filter: MetadataInput,
before: Optional[RunnableConfig] = None,
before: Optional[CheckpointConfig] = None,
) -> tuple[str, list[Any]]:
"""Return WHERE clause predicates for alist() given config, filter, before.
libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py
-928
View File
@@ -1,928 +0,0 @@
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,
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, 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(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 (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,
task_path,
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(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 (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,
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 (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,
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 (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.
task_path (str): Path of the task creating the writes.
"""
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id, task_path), self.loop
).result()
libs/checkpoint-postgres/langgraph/store/postgres/__init__.py
+1 -2
View File
@@ -1,4 +1,3 @@
from langgraph.store.postgres.aio import AsyncPostgresStore
from langgraph.store.postgres.base import PostgresStore
__all__ = ["AsyncPostgresStore", "PostgresStore"]
__all__ = ["PostgresStore"]
libs/checkpoint-postgres/langgraph/store/postgres/aio.py
-434
View File
@@ -1,434 +0,0 @@
import asyncio
import logging
from collections.abc import AsyncIterator, Iterable, Sequence
from contextlib import asynccontextmanager
from typing import Any, Callable, Optional, Union, cast
import orjson
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.rows import DictRow, dict_row
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.store.base import (
GetOp,
ListNamespacesOp,
Op,
PutOp,
Result,
SearchOp,
)
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.postgres.base import (
_PLACEHOLDER,
BasePostgresStore,
PoolConfig,
PostgresIndexConfig,
Row,
_decode_ns_bytes,
_ensure_index_config,
_group_ops,
_row_to_item,
_row_to_search_item,
)
logger = logging.getLogger(__name__)
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
"""Asynchronous Postgres-backed store with optional vector search using pgvector.
!!! example "Examples"
Basic setup and usage:
```python
from langgraph.store.postgres import AsyncPostgresStore
conn_string = "postgresql://user:pass@localhost:5432/dbname"
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
await store.setup() # Run migrations. Done once
# Store and retrieve data
await store.aput(("users", "123"), "prefs", {"theme": "dark"})
item = await store.aget(("users", "123"), "prefs")
```
Vector search using LangChain embeddings:
```python
from langchain.embeddings import init_embeddings
from langgraph.store.postgres import AsyncPostgresStore
conn_string = "postgresql://user:pass@localhost:5432/dbname"
async with AsyncPostgresStore.from_conn_string(
conn_string,
index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small"),
"fields": ["text"] # specify which fields to embed. Default is the whole serialized value
}
) as store:
await store.setup() # Run migrations. Done once
# Store documents
await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False) # don't index
# Search by similarity
results = await store.asearch(("docs",), "programming guides", limit=2)
```
Using connection pooling for better performance:
```python
from langgraph.store.postgres import AsyncPostgresStore, PoolConfig
conn_string = "postgresql://user:pass@localhost:5432/dbname"
async with AsyncPostgresStore.from_conn_string(
conn_string,
pool_config=PoolConfig(
min_size=5,
max_size=20
)
) as store:
await store.setup() # Run migrations. Done once
# Use store with connection pooling...
```
Warning:
Make sure to:
1. Call `setup()` before first use to create necessary tables and indexes
2. Have the pgvector extension available to use vector search
3. Use Python 3.10+ for async functionality
Note:
Semantic search is disabled by default. You can enable it by providing an `index` configuration
when creating the store. Without this configuration, all `index` arguments passed to
`put` or `aput` will have no effect.
"""
__slots__ = (
"_deserializer",
"pipe",
"lock",
"supports_pipeline",
"index_config",
"embeddings",
)
def __init__(
self,
conn: _ainternal.Conn,
*,
pipe: Optional[AsyncPipeline] = None,
deserializer: Optional[
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
] = None,
index: Optional[PostgresIndexConfig] = None,
) -> None:
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
raise ValueError(
"Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
)
super().__init__()
self._deserializer = deserializer
self.conn = conn
self.pipe = pipe
self.lock = asyncio.Lock()
self.loop = asyncio.get_running_loop()
self.supports_pipeline = Capabilities().has_pipeline()
self.index_config = index
if self.index_config:
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
else:
self.embeddings = None
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
async with _ainternal.get_connection(self.conn) as conn:
if self.pipe:
async with self.pipe:
await self._execute_batch(grouped_ops, results, conn)
else:
await self._execute_batch(grouped_ops, results, conn)
return results
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
*,
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
) -> AsyncIterator["AsyncPostgresStore"]:
"""Create a new AsyncPostgresStore instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
pipeline (bool): Whether to use AsyncPipeline (only for single connections)
pool_config (Optional[PoolConfig]): Configuration for the connection pool.
If provided, will create a connection pool and use it instead of a single connection.
This overrides the `pipeline` argument.
index (Optional[PostgresIndexConfig]): The embedding config.
Returns:
AsyncPostgresStore: A new AsyncPostgresStore instance.
"""
if pool_config is not None:
pc = pool_config.copy()
async with cast(
AsyncConnectionPool[AsyncConnection[DictRow]],
AsyncConnectionPool(
conn_string,
min_size=pc.pop("min_size", 1),
max_size=pc.pop("max_size", None),
kwargs={
"autocommit": True,
"prepare_threshold": 0,
"row_factory": dict_row,
**(pc.pop("kwargs", None) or {}),
},
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool, index=index)
else:
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, index=index)
else:
yield cls(conn=conn, index=index)
async def setup(self) -> None:
"""Set up the store 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 the store is used.
"""
async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
await cur.execute(
f"""
CREATE TABLE IF NOT EXISTS {table} (
v INTEGER PRIMARY KEY
)
"""
)
await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
row = cast(dict, await cur.fetchone())
if row is None:
version = -1
else:
version = row["v"]
return version
async with self._cursor() as cur:
version = await _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
await cur.execute(sql)
await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
if self.index_config:
version = await _get_version(cur, table="vector_migrations")
for v, migration in enumerate(
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
):
sql = migration.sql
if migration.params:
params = {
k: v(self) if v is not None and callable(v) else v
for k, v in migration.params.items()
}
sql = sql % params
await cur.execute(sql)
await cur.execute(
"INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
)
async def _execute_batch(
self,
grouped_ops: dict,
results: list[Result],
conn: AsyncConnection[DictRow],
) -> None:
async with self._cursor(pipeline=True) as cur:
if GetOp in grouped_ops:
await self._batch_get_ops(
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]),
results,
cur,
)
if SearchOp in grouped_ops:
await self._batch_search_ops(
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
results,
cur,
)
if ListNamespacesOp in grouped_ops:
await self._batch_list_namespaces_ops(
cast(
Sequence[tuple[int, ListNamespacesOp]],
grouped_ops[ListNamespacesOp],
),
results,
cur,
)
if PutOp in grouped_ops:
await self._batch_put_ops(
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]),
cur,
)
async def _batch_get_ops(
self,
get_ops: Sequence[tuple[int, GetOp]],
results: list[Result],
cur: AsyncCursor[DictRow],
) -> None:
for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
await cur.execute(query, params)
rows = cast(list[Row], await cur.fetchall())
key_to_row = {row["key"]: row for row in rows}
for idx, key in items:
row = key_to_row.get(key)
if row:
results[idx] = _row_to_item(
namespace, row, loader=self._deserializer
)
else:
results[idx] = None
async def _batch_put_ops(
self,
put_ops: Sequence[tuple[int, PutOp]],
cur: AsyncCursor[DictRow],
) -> None:
queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
if embedding_request:
if self.embeddings is None:
# Should not get here since the embedding config is required
# to return an embedding_request above
raise ValueError(
"Embedding configuration is required for vector operations "
f"(for semantic search). "
f"Please provide an EmbeddingConfig when initializing the {self.__class__.__name__}."
)
query, txt_params = embedding_request
vectors = await self.embeddings.aembed_documents(
[param[-1] for param in txt_params]
)
queries.append(
(
query,
[
p
for (ns, k, pathname, _), vector in zip(txt_params, vectors)
for p in (ns, k, pathname, vector)
],
)
)
for query, params in queries:
await cur.execute(query, params)
async def _batch_search_ops(
self,
search_ops: Sequence[tuple[int, SearchOp]],
results: list[Result],
cur: AsyncCursor[DictRow],
) -> None:
queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
if embedding_requests and self.embeddings:
vectors = await self.embeddings.aembed_documents(
[query for _, query in embedding_requests]
)
for (idx, _), vector in zip(embedding_requests, vectors):
_paramslist = queries[idx][1]
for i in range(len(_paramslist)):
if _paramslist[i] is _PLACEHOLDER:
_paramslist[i] = vector
for (idx, _), (query, params) in zip(search_ops, queries):
await cur.execute(query, params)
rows = cast(list[Row], await cur.fetchall())
items = [
_row_to_search_item(
_decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
)
for row in rows
]
results[idx] = items
async def _batch_list_namespaces_ops(
self,
list_ops: Sequence[tuple[int, ListNamespacesOp]],
results: list[Result],
cur: AsyncCursor[DictRow],
) -> None:
queries = self._get_batch_list_namespaces_queries(list_ops)
for (query, params), (idx, _) in zip(queries, list_ops):
await cur.execute(query, params)
rows = cast(list[dict], await cur.fetchall())
namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
results[idx] = namespaces
@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 PostgresStore 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:
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) as cur,
):
yield cur
libs/checkpoint-postgres/langgraph/store/postgres/base.py
+2 -6
View File
@@ -7,7 +7,6 @@ from collections.abc import Iterable, Iterator, Sequence
from contextlib import contextmanager
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
Callable,
Generic,
@@ -26,10 +25,10 @@ from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
from typing_extensions import TypedDict
from langgraph.checkpoint.postgres import _ainternal as _ainternal
from langgraph.checkpoint.postgres import _internal as _pg_internal
from langgraph.store.base import (
BaseStore,
Embeddings,
GetOp,
IndexConfig,
Item,
@@ -44,9 +43,6 @@ from langgraph.store.base import (
tokenize_path,
)
if TYPE_CHECKING:
from langchain_core.embeddings import Embeddings
logger = logging.getLogger(__name__)
@@ -127,7 +123,7 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vec
]
C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
C = TypeVar("C", bound=_pg_internal.Conn)
class PoolConfig(TypedDict, total=False):