Merge pull request #1360 from langchain-ai/nc/14auug/debug-checkpoint-full

Add missing attributes to checkpoint event in debug output

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],
         },
     }
 

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
 

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": [],
                 },
             },
         ]

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": [],
                 },
             },
         ]