From 66b9a7dee7e807e4ea32c0f3888ea55aceaa9831 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 15 Nov 2024 16:56:17 -0800 Subject: [PATCH 1/5] lib: When copying checkpoint, make it a child of the parent --- libs/langgraph/langgraph/pregel/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d2d693bef..f49a2207a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -961,7 +961,7 @@ class Pregel(PregelProtocol): next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = checkpointer.put( - checkpoint_config, + saved.parent_config or saved.config if saved else checkpoint_config, next_checkpoint, { **checkpoint_metadata, @@ -1222,7 +1222,7 @@ class Pregel(PregelProtocol): next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = await checkpointer.aput( - checkpoint_config, + saved.parent_config or saved.config if saved else checkpoint_config, next_checkpoint, { **checkpoint_metadata, From 36e49eb1903b68e846e9766ee2b30f2878bc4501 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 15 Nov 2024 17:04:40 -0800 Subject: [PATCH 2/5] Add distinct source --- libs/checkpoint/langgraph/checkpoint/base/__init__.py | 3 ++- libs/langgraph/langgraph/pregel/__init__.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index a3ad83b1a..6805ada0e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -39,12 +39,13 @@ PendingWrite = Tuple[str, str, Any] class CheckpointMetadata(TypedDict, total=False): """Metadata associated with a checkpoint.""" - source: Literal["input", "loop", "update"] + source: Literal["input", "loop", "update", "fork"] """The source of the checkpoint. - "input": The checkpoint was created from an input to invoke/stream/batch. - "loop": The checkpoint was created from inside the pregel loop. - "update": The checkpoint was created from a manual state update. + - "fork": The checkpoint was created as a copy of another checkpoint. """ step: int """The step number of the checkpoint. diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f49a2207a..b6bfa895a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -946,7 +946,7 @@ class Pregel(PregelProtocol): create_checkpoint(checkpoint, None, step), { **checkpoint_metadata, - "source": "update", + "source": "fork", "step": step + 1, "writes": {}, "parents": saved.metadata.get("parents", {}) if saved else {}, @@ -965,7 +965,7 @@ class Pregel(PregelProtocol): next_checkpoint, { **checkpoint_metadata, - "source": "update", + "source": "fork", "step": step + 1, "writes": {}, "parents": saved.metadata.get("parents", {}) if saved else {}, From 973ad76a587f0831a2e936620934a8db59d73d2e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 15 Nov 2024 17:05:18 -0800 Subject: [PATCH 3/5] Fix --- libs/langgraph/langgraph/pregel/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index b6bfa895a..e624d6b46 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -946,7 +946,7 @@ class Pregel(PregelProtocol): create_checkpoint(checkpoint, None, step), { **checkpoint_metadata, - "source": "fork", + "source": "update", "step": step + 1, "writes": {}, "parents": saved.metadata.get("parents", {}) if saved else {}, @@ -1226,7 +1226,7 @@ class Pregel(PregelProtocol): next_checkpoint, { **checkpoint_metadata, - "source": "update", + "source": "fork", "step": step + 1, "writes": {}, "parents": saved.metadata.get("parents", {}) if saved else {}, From d99dc7d81bdfb446aa85114c52fe1194af677625 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 15 Nov 2024 17:23:00 -0800 Subject: [PATCH 4/5] Fix missing writes --- libs/langgraph/langgraph/pregel/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index e624d6b46..ff081111c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -967,7 +967,6 @@ class Pregel(PregelProtocol): **checkpoint_metadata, "source": "fork", "step": step + 1, - "writes": {}, "parents": saved.metadata.get("parents", {}) if saved else {}, }, {}, @@ -1228,7 +1227,6 @@ class Pregel(PregelProtocol): **checkpoint_metadata, "source": "fork", "step": step + 1, - "writes": {}, "parents": saved.metadata.get("parents", {}) if saved else {}, }, {}, From 7c11325e2313c6fb738c96778911d3894a3c2181 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 15 Nov 2024 17:33:01 -0800 Subject: [PATCH 5/5] Separate --- libs/langgraph/langgraph/pregel/__init__.py | 36 +++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index ff081111c..59d2736a0 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -958,6 +958,24 @@ class Pregel(PregelProtocol): ) # no values, copy checkpoint if values is None and as_node is None: + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = checkpointer.put( + checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + if values is None and as_node == "__copy__": next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = checkpointer.put( @@ -1218,6 +1236,24 @@ class Pregel(PregelProtocol): ) # no values, copy checkpoint if values is None and as_node is None: + next_checkpoint = create_checkpoint(checkpoint, None, step) + # copy checkpoint + next_config = await checkpointer.aput( + checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "update", + "step": step + 1, + "writes": {}, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) + if values is None and as_node == "__copy__": next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = await checkpointer.aput(