diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 982182842..7e8f9546f 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -52,6 +52,7 @@ class CheckpointPayload(TypedDict): metadata: CheckpointMetadata values: dict[str, Any] next: list[str] + parent_config: Optional[RunnableConfig] tasks: list[CheckpointTask] @@ -136,6 +137,7 @@ def map_debug_checkpoint( checkpoint: Checkpoint, tasks: Iterable[PregelExecutableTask], pending_writes: list[PendingWrite], + parent_config: Optional[RunnableConfig], ) -> Iterator[DebugOutputCheckpoint]: """Produce "checkpoint" events for stream_mode=debug.""" yield { @@ -144,6 +146,7 @@ def map_debug_checkpoint( "step": step, "payload": { "config": config, + "parent_config": parent_config, "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 6353680c3..362756f32 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -176,6 +176,7 @@ class PregelLoop: checkpoint_metadata: CheckpointMetadata checkpoint_pending_writes: List[PendingWrite] checkpoint_previous_versions: dict[str, Union[str, float, int]] + prev_checkpoint_config: Optional[RunnableConfig] step: int stop: int @@ -249,6 +250,13 @@ class PregelLoop: if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) else () ) + self.prev_checkpoint_config = ( + self.checkpoint_config + if self.checkpoint_config + and CONF in self.checkpoint_config + and CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF] + else None + ) def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: """Put writes for a task, to be read by the next tick.""" @@ -386,6 +394,7 @@ class PregelLoop: self.checkpoint, self.tasks.values(), self.checkpoint_pending_writes, + self.prev_checkpoint_config, ) # if no more tasks, we're done @@ -535,6 +544,12 @@ 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_config = ( + self.checkpoint_config + if CONFIG_KEY_CHECKPOINT_ID in self.checkpoint_config[CONF] + and self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID] + else None + ) self.checkpoint_metadata = metadata self.checkpoint_config = { **self.checkpoint_config, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 636f450ee..d2943d068 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7019,6 +7019,7 @@ def test_branch_then( "step": -1, "writes": {"__start__": {"my_key": "value", "market": "DE"}}, }, + "parent_config": None, "next": ["__start__"], "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, @@ -7049,6 +7050,17 @@ def test_branch_then( "step": 0, "writes": None, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["prepare"], "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, @@ -7102,6 +7114,17 @@ def test_branch_then( "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["tool_two_slow"], "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], }, @@ -7155,6 +7178,17 @@ def test_branch_then( "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["finish"], "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, @@ -7208,6 +7242,17 @@ def test_branch_then( "step": 3, "writes": {"finish": {"my_key": " finished"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": [], "tasks": [], }, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 6a7252391..5ad973ce4 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5651,6 +5651,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": -1, "writes": {"__start__": {"my_key": "value", "market": "DE"}}, }, + "parent_config": None, "next": ["__start__"], "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, @@ -5681,6 +5682,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": 0, "writes": None, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["prepare"], "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, @@ -5734,6 +5746,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["tool_two_slow"], "tasks": [ {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} @@ -5789,6 +5812,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["finish"], "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, @@ -5842,6 +5876,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": 3, "writes": {"finish": {"my_key": " finished"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": [], "tasks": [], }, @@ -5888,6 +5933,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": -1, "writes": {"__start__": {"my_key": "value", "market": "DE"}}, }, + "parent_config": None, "next": ["__start__"], "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, @@ -5918,6 +5964,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": 0, "writes": None, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["prepare"], "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, @@ -5971,6 +6028,17 @@ async def test_branch_then(checkpointer_name: str) -> None: "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, "next": ["tool_two_slow"], "tasks": [ {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}