Lint

Author: nfcampos

libs/langgraph/langgraph/pregel/__init__.py

@@ -2236,8 +2236,6 @@ def _defaults(
             stream_mode = ["values"]
         elif stream_mode is None:
             stream_mode = self.stream_mode
-        elif stream_mode == "debug":
-            stream_mode = ["checkpoints", "tasks"]
         if not isinstance(stream_mode, list):
             stream_mode = [stream_mode]
         if self.checkpointer is False:

libs/langgraph/langgraph/pregel/debug.py

@@ -3,7 +3,6 @@
 from collections import defaultdict
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from dataclasses import asdict
-from datetime import datetime, timezone
 from pprint import pformat
 from typing import (
     Any,
@@ -17,7 +16,7 @@
 from typing_extensions import TypedDict
 
 from langgraph.channels.base import BaseChannel
-from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite
+from langgraph.checkpoint.base import CheckpointMetadata, PendingWrite
 from langgraph.constants import (
     CONF,
     CONFIG_KEY_CHECKPOINT_NS,
@@ -92,30 +91,21 @@ class DebugOutputCheckpoint(DebugOutputBase):
 TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
 
 
-def map_debug_tasks(
-    step: int, tasks: Iterable[PregelExecutableTask]
-) -> Iterator[DebugOutputTask]:
+def map_debug_tasks(tasks: Iterable[PregelExecutableTask]) -> Iterator[DebugOutputTask]:
     """Produce "task" events for stream_mode=debug."""
-    ts = datetime.now(timezone.utc).isoformat()
     for task in tasks:
         if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):
             continue
 
         yield {
-            "type": "task",
-            "timestamp": ts,
-            "step": step,
-            "payload": {
-                "id": task.id,
-                "name": task.name,
-                "input": task.input,
-                "triggers": task.triggers,
-            },
+            "id": task.id,
+            "name": task.name,
+            "input": task.input,
+            "triggers": task.triggers,
         }
 
 
 def map_debug_task_results(
-    step: int,
     task_tup: tuple[PregelExecutableTask, Sequence[tuple[str, Any]]],
     stream_keys: str | Sequence[str],
 ) -> Iterator[DebugOutputTaskResult]:
@@ -125,23 +115,16 @@ def map_debug_task_results(
     )
     task, writes = task_tup
     yield {
-        "type": "task_result",
-        "timestamp": datetime.now(timezone.utc).isoformat(),
-        "step": step,
-        "payload": {
-            "id": task.id,
-            "name": task.name,
-            "error": next((w[1] for w in writes if w[0] == ERROR), None),
-            "result": [
-                w for w in writes if w[0] in stream_channels_list or w[0] == RETURN
-            ],
-            "interrupts": [
-                asdict(v)
-                for w in writes
-                if w[0] == INTERRUPT
-                for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
-            ],
-        },
+        "id": task.id,
+        "name": task.name,
+        "error": next((w[1] for w in writes if w[0] == ERROR), None),
+        "result": [w for w in writes if w[0] in stream_channels_list or w[0] == RETURN],
+        "interrupts": [
+            asdict(v)
+            for w in writes
+            if w[0] == INTERRUPT
+            for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
+        ],
     }
 
 
@@ -159,12 +142,10 @@ def rm_pregel_keys(config: RunnableConfig | None) -> RunnableConfig | None:
 
 
 def map_debug_checkpoint(
-    step: int,
     config: RunnableConfig,
     channels: Mapping[str, BaseChannel],
     stream_channels: str | Sequence[str],
     metadata: CheckpointMetadata,
-    checkpoint: Checkpoint,
     tasks: Iterable[PregelExecutableTask],
     pending_writes: list[PendingWrite],
     parent_config: RunnableConfig | None,
@@ -193,42 +174,35 @@ def map_debug_checkpoint(
         }
 
     yield {
-        "type": "checkpoint",
-        "timestamp": checkpoint["ts"],
-        "step": step,
-        "payload": {
-            "config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),
-            "parent_config": rm_pregel_keys(
-                patch_checkpoint_map(parent_config, metadata)
-            ),
-            "values": read_channels(channels, stream_channels),
-            "metadata": metadata,
-            "next": [t.name for t in tasks],
-            "tasks": [
-                {
-                    "id": t.id,
-                    "name": t.name,
-                    "error": t.error,
-                    "state": t.state,
-                }
-                if t.error
-                else {
-                    "id": t.id,
-                    "name": t.name,
-                    "result": t.result,
-                    "interrupts": tuple(asdict(i) for i in t.interrupts),
-                    "state": t.state,
-                }
-                if t.result
-                else {
-                    "id": t.id,
-                    "name": t.name,
-                    "interrupts": tuple(asdict(i) for i in t.interrupts),
-                    "state": t.state,
-                }
-                for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)
-            ],
-        },
+        "config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),
+        "parent_config": rm_pregel_keys(patch_checkpoint_map(parent_config, metadata)),
+        "values": read_channels(channels, stream_channels),
+        "metadata": metadata,
+        "next": [t.name for t in tasks],
+        "tasks": [
+            {
+                "id": t.id,
+                "name": t.name,
+                "error": t.error,
+                "state": t.state,
+            }
+            if t.error
+            else {
+                "id": t.id,
+                "name": t.name,
+                "result": t.result,
+                "interrupts": tuple(asdict(i) for i in t.interrupts),
+                "state": t.state,
+            }
+            if t.result
+            else {
+                "id": t.id,
+                "name": t.name,
+                "interrupts": tuple(asdict(i) for i in t.interrupts),
+                "state": t.state,
+            }
+            for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)
+        ],
     }
 
 

libs/langgraph/langgraph/pregel/loop.py

@@ -11,6 +11,7 @@
     AsyncExitStack,
     ExitStack,
 )
+from datetime import datetime, timezone
 from inspect import signature
 from types import TracebackType
 from typing import (
@@ -423,7 +424,7 @@ def accept_push(
             ),
         ):
             # produce debug output
-            self._emit("tasks", map_debug_tasks, self.step, [pushed])
+            self._emit("tasks", map_debug_tasks, [pushed])
             # debug flag
             if self.debug:
                 print_step_tasks(self.step, [pushed])
@@ -475,7 +476,6 @@ def tick(self) -> bool:
             self._emit(
                 "checkpoints",
                 map_debug_checkpoint,
-                self.step - 1,  # printing checkpoint for previous step
                 {
                     **self.checkpoint_config,
                     CONF: {
@@ -486,7 +486,6 @@ def tick(self) -> bool:
                 self.channels,
                 self.stream_keys,
                 self.checkpoint_metadata,
-                self.checkpoint,
                 self.tasks.values(),
                 self.checkpoint_pending_writes,
                 self.prev_checkpoint_config,
@@ -510,7 +509,7 @@ def tick(self) -> bool:
             raise GraphInterrupt()
 
         # produce debug output
-        self._emit("tasks", map_debug_tasks, self.step, self.tasks.values())
+        self._emit("tasks", map_debug_tasks, self.tasks.values())
 
         # debug flag
         if self.debug:
@@ -842,10 +841,32 @@ def _emit(
     ) -> None:
         if self.stream is None:
             return
-        if mode not in self.stream.modes:
+        debug_remap = mode in ("checkpoints", "tasks") and "debug" in self.stream.modes
+        if mode not in self.stream.modes and not debug_remap:
             return
         for v in values(*args, **kwargs):
-            self.stream((self.checkpoint_ns, mode, v))
+            if mode in self.stream.modes:
+                self.stream((self.checkpoint_ns, mode, v))
+            # "debug" mode is "checkpoints" or "tasks" with a wrapper dict
+            if debug_remap:
+                self.stream(
+                    (
+                        self.checkpoint_ns,
+                        "debug",
+                        {
+                            "step": self.step - 1
+                            if mode == "checkpoints"
+                            else self.step,
+                            "timestamp": datetime.now(timezone.utc).isoformat(),
+                            "type": "checkpoint"
+                            if mode == "checkpoints"
+                            else "task_result"
+                            if "result" in v
+                            else "task",
+                            "payload": v,
+                        },
+                    )
+                )
 
     def output_writes(
         self, task_id: str, writes: WritesT, *, cached: bool = False
@@ -888,7 +909,6 @@ def output_writes(
                 self._emit(
                     "tasks",
                     map_debug_task_results,
-                    self.step,
                     (task, writes),
                     self.stream_keys,
                 )

libs/langgraph/langgraph/types.py

@@ -46,7 +46,9 @@ class ToolOutputMixin:  # type: ignore[no-redef]
 - False disables checkpointing, even if the parent graph has a checkpointer.
 - None inherits checkpointer from the parent graph."""
 
-StreamMode = Literal["values", "updates", "checkpoints", "tasks", "messages", "custom"]
+StreamMode = Literal[
+    "values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"
+]
 """How the stream method should emit outputs.
 
 - `"values"`: Emit all values in the state after each step, including interrupts.
@@ -57,6 +59,7 @@ class ToolOutputMixin:  # type: ignore[no-redef]
 - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks.
 - `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by get_state().
 - `"tasks"`: Emit events when tasks start and finish, including their results and errors.
+- `"debug"`: Emit "checlkpoints" and "tasks" events, for debugging purposes.
 """
 
 StreamWriter = Callable[[Any], None]

libs/sdk-py/langgraph_sdk/schema.py

@@ -36,14 +36,24 @@
 """
 
 StreamMode = Literal[
-    "values", "messages", "updates", "events", "debug", "custom", "messages-tuple"
+    "values",
+    "messages",
+    "updates",
+    "events",
+    "tasks",
+    "checkpoints",
+    "debug",
+    "custom",
+    "messages-tuple",
 ]
 """
 Defines the mode of streaming:
 - "values": Stream only the values.
 - "messages": Stream complete messages.
 - "updates": Stream updates to the state.
 - "events": Stream events occurring during execution.
+- "checkpoints": Stream checkpoints as they are created.
+- "tasks": Stream task start and finish events.
 - "debug": Stream detailed debug information.
 - "custom": Stream custom events.
 """