From 2c8fc3c9c67085065c733910b029427f28271376 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 14 Aug 2024 17:48:35 -0700 Subject: [PATCH 1/2] Add missing attributes to checkpoint event in debug output --- libs/langgraph/langgraph/pregel/debug.py | 13 ++++++++---- libs/langgraph/langgraph/pregel/loop.py | 27 ++++++++++++++---------- 2 files changed, 25 insertions(+), 15 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 9e152cf87..4aed748ae 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -9,8 +9,8 @@ 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.constants import TAG_HIDDEN +from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite +from langgraph.constants import INTERRUPT, TAG_HIDDEN from langgraph.pregel.io import read_channels from langgraph.pregel.types import PregelExecutableTask @@ -118,16 +118,21 @@ 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": { + "id": checkpoint["id"], "config": config, "values": read_channels(channels, stream_channels), "metadata": metadata, + "next": [t.name for t in tasks], + "interrupts": [v for tid, n, v in pending_writes or [] if n == INTERRUPT], }, } diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 0508cd942..f65b8174f 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, + 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 From c701fef6345e0a00a2be7a042592d9a95f5d88c0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 15 Aug 2024 09:40:53 -0700 Subject: [PATCH 2/2] Lint --- libs/langgraph/langgraph/pregel/debug.py | 4 +--- libs/langgraph/langgraph/pregel/loop.py | 2 +- libs/langgraph/tests/test_pregel.py | 5 +++++ libs/langgraph/tests/test_pregel_async.py | 5 +++++ 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 4aed748ae..3dbb46675 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -10,7 +10,7 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite -from langgraph.constants import INTERRUPT, TAG_HIDDEN +from langgraph.constants import TAG_HIDDEN from langgraph.pregel.io import read_channels from langgraph.pregel.types import PregelExecutableTask @@ -127,12 +127,10 @@ def map_debug_checkpoint( "timestamp": checkpoint["ts"], "step": step, "payload": { - "id": checkpoint["id"], "config": config, "values": read_channels(channels, stream_channels), "metadata": metadata, "next": [t.name for t in tasks], - "interrupts": [v for tid, n, v in pending_writes or [] if n == INTERRUPT], }, } diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f65b8174f..cd4bff66a 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -233,7 +233,7 @@ class PregelLoop: self.stream.extend( ("debug", v) for v in map_debug_checkpoint( - self.step, + self.step - 1, # printing checkpoint for previous step self.checkpoint_config, self.channels, self.graph.stream_channels_asis, 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": [], }, }, ]