diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 9e152cf87..3dbb46675 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -9,7 +9,7 @@ from langchain_core.runnables.config import RunnableConfig from langchain_core.utils.input import get_bolded_text, get_colored_text from langgraph.channels.base import BaseChannel -from langgraph.checkpoint.base import CheckpointMetadata +from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite from langgraph.constants import TAG_HIDDEN from langgraph.pregel.io import read_channels from langgraph.pregel.types import PregelExecutableTask @@ -118,16 +118,19 @@ def map_debug_checkpoint( channels: Mapping[str, BaseChannel], stream_channels: Union[str, Sequence[str]], metadata: CheckpointMetadata, + checkpoint: Checkpoint, + tasks: list[PregelExecutableTask], + pending_writes: list[PendingWrite], ) -> Iterator[DebugOutputCheckpoint]: - ts = datetime.now(timezone.utc).isoformat() yield { "type": "checkpoint", - "timestamp": ts, + "timestamp": checkpoint["ts"], "step": step, "payload": { "config": 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 0508cd942..cd4bff66a 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -228,6 +228,22 @@ class PregelLoop: is_resuming=self.input is INPUT_RESUMING, ) + # produce debug output + if self._checkpointer_put_after_previous is not None: + self.stream.extend( + ("debug", v) + for v in map_debug_checkpoint( + self.step - 1, # printing checkpoint for previous step + self.checkpoint_config, + self.channels, + self.graph.stream_channels_asis, + self.checkpoint_metadata, + self.checkpoint, + self.tasks, + self.checkpoint_pending_writes, + ) + ) + # if no more tasks, we're done if not self.tasks: self.status = "done" @@ -361,17 +377,6 @@ class PregelLoop: "checkpoint_id": self.checkpoint["id"], }, } - # produce debug output - self.stream.extend( - ("debug", v) - for v in map_debug_checkpoint( - self.step, - self.checkpoint_config, - self.channels, - self.graph.stream_channels_asis, - self.checkpoint_metadata, - ) - ) # increment step self.step += 1 diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index af7f0920a..d6a83e783 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6260,6 +6260,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "step": -1, "writes": {"my_key": "value", "market": "DE"}, }, + "next": ["__start__"], }, }, { @@ -6287,6 +6288,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "step": 0, "writes": None, }, + "next": ["prepare"], }, }, { @@ -6335,6 +6337,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, }, + "next": ["tool_two_slow"], }, }, { @@ -6383,6 +6386,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, }, + "next": ["finish"], }, }, { @@ -6431,6 +6435,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "step": 3, "writes": {"finish": {"my_key": " finished"}}, }, + "next": [], }, }, ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 67b62941b..ef49aed62 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4885,6 +4885,7 @@ async def test_branch_then() -> None: "step": -1, "writes": {"my_key": "value", "market": "DE"}, }, + "next": ["__start__"], }, }, { @@ -4912,6 +4913,7 @@ async def test_branch_then() -> None: "step": 0, "writes": None, }, + "next": ["prepare"], }, }, { @@ -4960,6 +4962,7 @@ async def test_branch_then() -> None: "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, }, + "next": ["tool_two_slow"], }, }, { @@ -5008,6 +5011,7 @@ async def test_branch_then() -> None: "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, }, + "next": ["finish"], }, }, { @@ -5056,6 +5060,7 @@ async def test_branch_then() -> None: "step": 3, "writes": {"finish": {"my_key": " finished"}}, }, + "next": [], }, }, ]