diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index f70a6e8c4..1fc922761 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -32,6 +32,7 @@ from langgraph.constants import ( from langgraph.pregel.io import read_channels from langgraph.pregel.utils import find_subgraph_pregel from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot +from langgraph.utils.config import patch_checkpoint_map class TaskPayload(TypedDict): @@ -177,8 +178,9 @@ def map_debug_checkpoint( "timestamp": checkpoint["ts"], "step": step, "payload": { - "config": config, + "config": patch_checkpoint_map(config, metadata), "parent_config": parent_config, + # "parent_config": patch_checkpoint_map(parent_config, metadata), "values": read_channels(channels, stream_channels), "metadata": metadata, "next": [t.name for t in tasks], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index e472bdd1a..d4f3a52c3 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -101,7 +101,7 @@ from langgraph.pregel.read import PregelNode from langgraph.pregel.utils import get_new_channel_versions from langgraph.store.base import BaseStore from langgraph.types import All, PregelExecutableTask, StreamMode -from langgraph.utils.config import patch_checkpoint_map, patch_configurable +from langgraph.utils.config import patch_configurable V = TypeVar("V") P = ParamSpec("P") @@ -176,7 +176,6 @@ class PregelLoop: checkpoint_pending_writes: List[PendingWrite] checkpoint_previous_versions: dict[str, Union[str, float, int]] prev_checkpoint_config: Optional[RunnableConfig] - prev_checkpoint_metadata: Optional[CheckpointMetadata] step: int stop: int @@ -363,16 +362,14 @@ class PregelLoop: "debug", map_debug_checkpoint, self.step - 1, # printing checkpoint for previous step - patch_checkpoint_map(self.checkpoint_config, self.checkpoint_metadata), + self.checkpoint_config, self.channels, self.stream_keys, self.checkpoint_metadata, self.checkpoint, self.tasks.values(), self.checkpoint_pending_writes, - patch_checkpoint_map( - self.prev_checkpoint_config, self.prev_checkpoint_metadata - ), + self.prev_checkpoint_config, self.output_keys, ) @@ -505,7 +502,7 @@ class PregelLoop: ) def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: - # assign step + # assign step and parents metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) # debug flag @@ -521,7 +518,6 @@ class PregelLoop: self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step) # bail if no checkpointer if self._checkpointer_put_after_previous is not None: - self.prev_checkpoint_metadata = self.checkpoint_metadata self.checkpoint_metadata = metadata self.prev_checkpoint_config = ( diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index 3993df9b2..fe25b6d9a 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -36,9 +36,11 @@ def patch_configurable( def patch_checkpoint_map( - config: RunnableConfig, metadata: Optional[CheckpointMetadata] + config: Optional[RunnableConfig], metadata: Optional[CheckpointMetadata] ) -> RunnableConfig: - if parents := (metadata.get("parents") if metadata else None): + if config is None: + return config + elif parents := (metadata.get("parents") if metadata else None): conf = config[CONF] return patch_configurable( config, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a17cf4e99..febc85f24 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -11917,7 +11917,6 @@ def test_debug_nested_subgraphs(): clean_config["thread_id"] = config["configurable"]["thread_id"] clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"] clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] - clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"] if "checkpoint_map" in config["configurable"]: clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"]