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 d2d693bef..59d2736a0 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -975,6 +975,23 @@ class Pregel(PregelProtocol): 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( + saved.parent_config or saved.config if saved else checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "fork", + "step": step + 1, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) # apply pending writes, if not on specific checkpoint if ( CONFIG_KEY_CHECKPOINT_ID not in config[CONF] @@ -1236,6 +1253,23 @@ class Pregel(PregelProtocol): 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( + saved.parent_config or saved.config if saved else checkpoint_config, + next_checkpoint, + { + **checkpoint_metadata, + "source": "fork", + "step": step + 1, + "parents": saved.metadata.get("parents", {}) if saved else {}, + }, + {}, + ) + return patch_checkpoint_map( + next_config, saved.metadata if saved else None + ) # apply pending writes, if not on specific checkpoint if ( CONFIG_KEY_CHECKPOINT_ID not in config[CONF]