WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-24 16:42:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

latest

Browse Source
This commit is contained in:
Sydney Runkle
2026-04-22 08:48:53 -04:00
parent 43ecd9dc2d
commit 1099cece53
15 changed files with 475 additions and 614 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+4 -39
View File
@@ -32,7 +32,8 @@ Conn = _internal.Conn # For backward compatibility
class PostgresSaver(BasePostgresSaver):
"""Checkpointer that stores checkpoints in a Postgres database."""
lock: threading.Lock
supports_delta_channels: bool = True
lock: threading.RLock
def __init__(
self,
@@ -48,7 +49,7 @@ class PostgresSaver(BasePostgresSaver):
self.conn = conn
self.pipe = pipe
self.lock = threading.Lock()
self.lock = threading.RLock()
self.supports_pipeline = Capabilities().has_pipeline()
@classmethod
@@ -430,43 +431,6 @@ class PostgresSaver(BasePostgresSaver):
with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
def get_channel_blob(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Look up a channel blob by checkpoint ID + channel via checkpoint_blobs."""
with self._cursor() as cur:
cur.execute(
"""
SELECT cb.type, cb.blob
FROM checkpoint_blobs cb
WHERE cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (
SELECT checkpoint->'channel_versions'->>%s
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
)
""",
(
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
checkpoint_id,
),
)
row = cur.fetchone()
if row is None:
return NotImplemented
return self.serde.loads_typed((row["type"], row["blob"]))
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
Convert a database row into a CheckpointTuple object.
@@ -484,6 +448,7 @@ class PostgresSaver(BasePostgresSaver):
value["channel_values"],
thread_id=value["thread_id"],
checkpoint_ns=value["checkpoint_ns"],
checkpoint_id=value["checkpoint_id"],
cur=cur,
)
return CheckpointTuple(
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+81 -30
View File
@@ -13,6 +13,8 @@ from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
CheckpointTuple,
DeltaChainValue,
DeltaValue,
get_checkpoint_id,
get_serializable_checkpoint_metadata,
)
@@ -32,6 +34,7 @@ Conn = _ainternal.Conn # For backward compatibility
class AsyncPostgresSaver(BasePostgresSaver):
"""Asynchronous checkpointer that stores checkpoints in a Postgres database."""
supports_delta_channels: bool = True
lock: asyncio.Lock
def __init__(
@@ -391,42 +394,80 @@ class AsyncPostgresSaver(BasePostgresSaver):
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
async def aget_channel_blob(
async def _aload_delta_chain(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
) -> Any:
"""Async look up of a channel blob by checkpoint ID + channel name."""
async with self._cursor() as cur:
await cur.execute(
"""
SELECT cb.type, cb.blob
FROM checkpoint_blobs cb
WHERE cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (
SELECT checkpoint->'channel_versions'->>%s
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
)
""",
(
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
checkpoint_id,
),
cur: Any,
) -> DeltaChainValue:
"""Fetch the full delta chain for a channel in one recursive CTE query (async)."""
await cur.execute(
"""
WITH RECURSIVE chain AS (
SELECT
c.checkpoint_id,
c.parent_checkpoint_id,
cb.version,
cb.type,
cb.blob
FROM checkpoints c
JOIN checkpoint_blobs cb
ON cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
WHERE c.thread_id = %s
AND c.checkpoint_ns = %s
AND c.checkpoint_id = %s
UNION ALL
SELECT
c.checkpoint_id,
c.parent_checkpoint_id,
cb.version,
cb.type,
cb.blob
FROM chain prev
JOIN checkpoints c ON c.checkpoint_id = prev.parent_checkpoint_id
JOIN checkpoint_blobs cb
ON cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
WHERE prev.parent_checkpoint_id IS NOT NULL
AND prev.type = 'delta'
)
row = await cur.fetchone()
if row is None:
return NotImplemented
return self.serde.loads_typed((row["type"], row["blob"]))
SELECT DISTINCT ON (version) type, blob
FROM chain
ORDER BY version ASC
""",
(
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
checkpoint_id,
thread_id,
checkpoint_ns,
channel,
channel,
),
)
rows = await cur.fetchall()
base = None
deltas: list[list[Any]] = []
for row in rows:
blob = self.serde.loads_typed((row["type"], row["blob"]))
if isinstance(blob, DeltaValue):
deltas.append(blob.delta)
else:
base = blob
return DeltaChainValue(base=base, deltas=deltas)
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
"""
@@ -442,11 +483,21 @@ class AsyncPostgresSaver(BasePostgresSaver):
"""
thread_id = value["thread_id"]
checkpoint_ns = value["checkpoint_ns"]
checkpoint_id = value["checkpoint_id"]
blob_values = value["channel_values"]
channel_values: dict[str, Any] = {}
if blob_values:
channel_values = self._load_blobs(blob_values)
delta_channels = [
k.decode() for k, t, _ in blob_values if t.decode() == "delta"
]
if delta_channels:
async with self._cursor() as cur:
for channel in delta_channels:
channel_values[channel] = await self._aload_delta_chain(
thread_id, checkpoint_ns, checkpoint_id, channel, cur
)
return CheckpointTuple(
{
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+87 -1
View File
@@ -11,6 +11,8 @@ from langgraph.checkpoint.base import (
WRITES_IDX_MAP,
BaseCheckpointSaver,
ChannelVersions,
DeltaChainValue,
DeltaValue,
get_checkpoint_id,
)
from langgraph.checkpoint.serde.types import TASKS
@@ -190,18 +192,102 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
*,
thread_id: str = "",
checkpoint_ns: str = "",
checkpoint_id: str = "",
cur: Any = None,
) -> dict[str, Any]:
if not blob_values:
return {}
result: dict[str, Any] = {}
delta_channels: list[str] = []
for k, t, v in blob_values:
channel = k.decode()
type_tag = t.decode()
if type_tag != "empty":
if type_tag == "delta":
delta_channels.append(channel)
elif type_tag != "empty":
result[channel] = self.serde.loads_typed((type_tag, v))
if delta_channels and cur is not None and checkpoint_id:
for channel in delta_channels:
result[channel] = self._load_delta_chain(
thread_id, checkpoint_ns, checkpoint_id, channel, cur
)
return result
def _load_delta_chain(
self,
thread_id: str,
checkpoint_ns: str,
checkpoint_id: str,
channel: str,
cur: Any,
) -> DeltaChainValue:
"""Fetch the full delta chain for a channel in one recursive CTE query."""
cur.execute(
"""
WITH RECURSIVE chain AS (
SELECT
c.checkpoint_id,
c.parent_checkpoint_id,
cb.version,
cb.type,
cb.blob
FROM checkpoints c
JOIN checkpoint_blobs cb
ON cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
WHERE c.thread_id = %s
AND c.checkpoint_ns = %s
AND c.checkpoint_id = %s
UNION ALL
SELECT
c.checkpoint_id,
c.parent_checkpoint_id,
cb.version,
cb.type,
cb.blob
FROM chain prev
JOIN checkpoints c ON c.checkpoint_id = prev.parent_checkpoint_id
JOIN checkpoint_blobs cb
ON cb.thread_id = %s
AND cb.checkpoint_ns = %s
AND cb.channel = %s
AND cb.version = (c.checkpoint->'channel_versions'->>%s)::text
WHERE prev.parent_checkpoint_id IS NOT NULL
AND prev.type = 'delta'
)
SELECT DISTINCT ON (version) type, blob
FROM chain
ORDER BY version ASC
""",
(
thread_id,
checkpoint_ns,
channel,
channel,
thread_id,
checkpoint_ns,
checkpoint_id,
thread_id,
checkpoint_ns,
channel,
channel,
),
)
rows = cur.fetchall()
base = None
deltas: list[list[Any]] = []
for row in rows:
blob = self.serde.loads_typed((row["type"], row["blob"]))
if isinstance(blob, DeltaValue):
deltas.append(blob.delta)
else:
base = blob
return DeltaChainValue(base=base, deltas=deltas)
def _dump_blobs(
self,
thread_id: str,