Merge pull request #2434 from langchain-ai/nc/15nov/update-state-copy-parent

lib: When copying checkpoint, make it a child of the parent

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.

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]