mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 07:32:25 +02:00
fix(langgraph): fix PostgresSaver crashing when loading older checkpoints (#6162)
### Description https://github.com/langchain-ai/langgraph/issues/6137 and https://github.com/langchain-ai/langgraph/issues/5677 reported issues where older checkpoints read by AsyncPostgresSaver/PostgresSaver from `langgraph-checkpoint-postgres==2.0.19` fail to read channel values, throwing `NoneType object is not a mapping`. This was due to a bug in how `channel_values` is assembled: ```python "channel_values": { **value["checkpoint"].get("channel_values"), # <--- if channel_values doesn't exist (old checkpoint), **None errors **self._load_blobs(value["channel_values"]), }, ``` This bug was observed for checkpoints generated by `langgraph-checkpoint-postgres<=2.0.19`. Fixed by providing a fallback to `value["checkpoint"].get("channel_values")`: ```python **value["checkpoint"], "channel_values": { **( value["checkpoint"].get("channel_values") or {} ), # 'or {}' needed for backwards compat with v3 checkpoints and below, as v4 introduced channel_values key **self._load_blobs(value["channel_values"]), }, ``` ### Tests Added test for AsyncPostgresSaver and test for PostgresSaver, using monkeypatch to remove `channel_values` before CheckpointTuple is assembled in `_load_checkpoint_tuple`. ### Solves https://github.com/langchain-ai/langgraph/issues/6137 and https://github.com/langchain-ai/langgraph/issues/5677 --------- Co-authored-by: Shahrukh Shaik <144558473+shahrukh-shaik@users.noreply.github.com>
This commit is contained in:
co-authored by
Shahrukh Shaik
parent
8dc4465d05
commit
f0fced262a
@@ -450,7 +450,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**value["checkpoint"].get("channel_values"),
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
},
|
||||
|
||||
@@ -409,7 +409,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
{
|
||||
**value["checkpoint"],
|
||||
"channel_values": {
|
||||
**value["checkpoint"].get("channel_values"),
|
||||
**(value["checkpoint"].get("channel_values") or {}),
|
||||
**self._load_blobs(value["channel_values"]),
|
||||
},
|
||||
},
|
||||
|
||||
@@ -344,3 +344,34 @@ async def test_pending_sends_migration(saver_name: str) -> None:
|
||||
TASKS: ["send-1", "send-2", "send-3"]
|
||||
}
|
||||
assert TASKS in search_results[0].checkpoint["channel_versions"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
async def test_get_checkpoint_no_channel_values(
|
||||
monkeypatch, saver_name: str, test_data
|
||||
) -> None:
|
||||
"""Backwards compatibility test that verifies a checkpoint with no channel_values key can be retrieved without throwing an error."""
|
||||
async with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_ns": "",
|
||||
"__super_private_key": "super_private_value",
|
||||
},
|
||||
"metadata": {"run_id": "my_run_id"},
|
||||
}
|
||||
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
||||
await saver.aput(config, chkpnt, {}, {})
|
||||
|
||||
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
||||
|
||||
def patched_load_checkpoint_tuple(value):
|
||||
value["checkpoint"].pop("channel_values", None)
|
||||
return load_checkpoint_tuple(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
||||
)
|
||||
|
||||
checkpoint = await saver.aget_tuple(config)
|
||||
assert checkpoint.checkpoint["channel_values"] == {}
|
||||
|
||||
@@ -332,3 +332,33 @@ def test_pending_sends_migration(saver_name: str) -> None:
|
||||
TASKS: ["send-1", "send-2", "send-3"]
|
||||
}
|
||||
assert TASKS in search_results[0].checkpoint["channel_versions"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
||||
def test_get_checkpoint_no_channel_values(
|
||||
monkeypatch, saver_name: str, test_data
|
||||
) -> None:
|
||||
"""Backwards compatibility test that verifies a checkpoint with no channel_values key can be retrieved without throwing an error."""
|
||||
with _saver(saver_name) as saver:
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "thread-2",
|
||||
"checkpoint_ns": "",
|
||||
"__super_private_key": "super_private_value",
|
||||
},
|
||||
}
|
||||
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
||||
saver.put(config, chkpnt, {}, {})
|
||||
|
||||
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
||||
|
||||
def patched_load_checkpoint_tuple(value):
|
||||
value["checkpoint"].pop("channel_values", None)
|
||||
return load_checkpoint_tuple(value)
|
||||
|
||||
monkeypatch.setattr(
|
||||
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
||||
)
|
||||
|
||||
checkpoint = saver.get_tuple(config)
|
||||
assert checkpoint.checkpoint["channel_values"] == {}
|
||||
|
||||
Reference in New Issue
Block a user