mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 00:22:25 +02:00
feat(checkpoint/postgres): diff chain reconstruction in _load_blobs (sync)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
dba1987c9b
commit
d4e1efa1f6
@@ -430,6 +430,46 @@ class PostgresSaver(BasePostgresSaver):
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def _load_diff_chains(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
diff_channel_payloads: dict[str, dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
from langgraph.checkpoint.base import DiffChainValue
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
with self._cursor() as cur:
|
||||
for channel, current_payload in diff_channel_payloads.items():
|
||||
payloads: list[dict[str, Any]] = [current_payload]
|
||||
version_cursor: str | None = current_payload["p"]
|
||||
base: list[Any] | None = None
|
||||
|
||||
while version_cursor is not None:
|
||||
cur.execute(
|
||||
"SELECT type, blob FROM checkpoint_blobs "
|
||||
"WHERE thread_id = %s AND checkpoint_ns = %s "
|
||||
"AND channel = %s AND version = %s",
|
||||
(thread_id, checkpoint_ns, channel, version_cursor),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
break
|
||||
# row is a dict (dict_row factory): {"type": str, "blob": bytes}
|
||||
if row["type"] == "diff":
|
||||
payload = self.serde.loads_typed(("diff", row["blob"]))
|
||||
payloads.append(payload)
|
||||
version_cursor = payload["p"]
|
||||
else:
|
||||
base = self.serde.loads_typed((row["type"], row["blob"]))
|
||||
break
|
||||
|
||||
payloads.reverse()
|
||||
result[channel] = DiffChainValue(
|
||||
base=base, deltas=[p["d"] for p in payloads]
|
||||
)
|
||||
return result
|
||||
|
||||
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
|
||||
"""
|
||||
Convert a database row into a CheckpointTuple object.
|
||||
@@ -454,7 +494,11 @@ class PostgresSaver(BasePostgresSaver):
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
**self._load_blobs(
|
||||
value["channel_values"],
|
||||
thread_id=value["thread_id"],
|
||||
checkpoint_ns=value["checkpoint_ns"],
|
||||
),
|
||||
},
|
||||
},
|
||||
value["metadata"],
|
||||
|
||||
@@ -185,15 +185,43 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
)
|
||||
|
||||
def _load_blobs(
|
||||
self, blob_values: list[tuple[bytes, bytes, bytes]]
|
||||
self,
|
||||
blob_values: list[tuple[bytes, bytes, bytes]],
|
||||
*,
|
||||
thread_id: str = "",
|
||||
checkpoint_ns: str = "",
|
||||
) -> dict[str, Any]:
|
||||
if not blob_values:
|
||||
return {}
|
||||
return {
|
||||
k.decode(): self.serde.loads_typed((t.decode(), v))
|
||||
for k, t, v in blob_values
|
||||
if t.decode() != "empty"
|
||||
}
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
diff_channel_payloads: dict[str, dict[str, Any]] = {}
|
||||
|
||||
for k, t, v in blob_values:
|
||||
channel = k.decode()
|
||||
type_tag = t.decode()
|
||||
if type_tag == "diff":
|
||||
diff_channel_payloads[channel] = self.serde.loads_typed(
|
||||
(type_tag, v)
|
||||
)
|
||||
elif type_tag != "empty":
|
||||
result[channel] = self.serde.loads_typed((type_tag, v))
|
||||
|
||||
if diff_channel_payloads:
|
||||
result.update(
|
||||
self._load_diff_chains(thread_id, checkpoint_ns, diff_channel_payloads)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def _load_diff_chains(
|
||||
self,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
diff_channel_payloads: dict[str, dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Override in sync/async subclasses. Resolves diff-chain blobs to DiffChainValue."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
|
||||
@@ -49,6 +49,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
def _is_diff_delta(obj: Any) -> bool:
|
||||
from langgraph.checkpoint.base import DiffDelta # lazy import avoids circular dep
|
||||
|
||||
return isinstance(obj, DiffDelta)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user