From 25a59447c1cbe9fdb448081a83ee3b3752fb7be3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 16 Jun 2025 08:47:45 -0700 Subject: [PATCH 1/2] Introduce "tasks" and "checkpoints" stream modes - These are split out of "debug" stream mode, which is now an alias for ["tasks", "checkpoints"] --- libs/langgraph/langgraph/pregel/__init__.py | 5 ++++- libs/langgraph/langgraph/pregel/loop.py | 11 ++++++----- libs/langgraph/langgraph/types.py | 5 +++-- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3ee6420fa..84eef0abe 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2236,6 +2236,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou 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: @@ -2298,7 +2300,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou - `"custom"`: Emit custom data from inside nodes or tasks using `StreamWriter`. - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks. Will be emitted as 2-tuples `(LLM token, metadata)`. - - `"debug"`: Emit debug events with as much information as possible for each step. + - `"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. You can pass a list as the `stream_mode` parameter to stream multiple modes at once. The streamed outputs will be tuples of `(mode, data)`. diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index ee4382557..9043d6041 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -119,6 +119,7 @@ from langgraph.types import ( PregelScratchpad, RetryPolicy, StreamChunk, + StreamMode, StreamProtocol, ) from langgraph.utils.config import patch_configurable @@ -422,7 +423,7 @@ class PregelLoop: ), ): # produce debug output - self._emit("debug", map_debug_tasks, self.step, [pushed]) + self._emit("tasks", map_debug_tasks, self.step, [pushed]) # debug flag if self.debug: print_step_tasks(self.step, [pushed]) @@ -472,7 +473,7 @@ class PregelLoop: # produce debug output if self._checkpointer_put_after_previous is not None: self._emit( - "debug", + "checkpoints", map_debug_checkpoint, self.step - 1, # printing checkpoint for previous step { @@ -509,7 +510,7 @@ class PregelLoop: raise GraphInterrupt() # produce debug output - self._emit("debug", map_debug_tasks, self.step, self.tasks.values()) + self._emit("tasks", map_debug_tasks, self.step, self.tasks.values()) # debug flag if self.debug: @@ -834,7 +835,7 @@ class PregelLoop: def _emit( self, - mode: str, + mode: StreamMode, values: Callable[P, Iterator[Any]], *args: P.args, **kwargs: P.kwargs, @@ -885,7 +886,7 @@ class PregelLoop: ) if not cached: self._emit( - "debug", + "tasks", map_debug_task_results, self.step, (task, writes), diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3e907e3eb..a5004ab2d 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -46,7 +46,7 @@ Checkpointer = Union[None, bool, BaseCheckpointSaver] - False disables checkpointing, even if the parent graph has a checkpointer. - None inherits checkpointer from the parent graph.""" -StreamMode = Literal["values", "updates", "debug", "messages", "custom"] +StreamMode = Literal["values", "updates", "checkpoints", "tasks", "messages", "custom"] """How the stream method should emit outputs. - `"values"`: Emit all values in the state after each step, including interrupts. @@ -55,7 +55,8 @@ StreamMode = Literal["values", "updates", "debug", "messages", "custom"] If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately. - `"custom"`: Emit custom data using from inside nodes or tasks using `StreamWriter`. - `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks. -- `"debug"`: Emit debug events with as much information as possible for each step. +- `"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. """ StreamWriter = Callable[[Any], None] From 417103066bde3680530e83b5dda0cc88b3daae77 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 16 Jun 2025 08:51:18 -0700 Subject: [PATCH 2/2] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 2 - libs/langgraph/langgraph/pregel/debug.py | 149 +++++++------------- libs/langgraph/langgraph/pregel/loop.py | 34 ++++- libs/langgraph/langgraph/types.py | 5 +- libs/sdk-py/langgraph_sdk/schema.py | 12 +- 5 files changed, 90 insertions(+), 112 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 84eef0abe..3946ac746 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2236,8 +2236,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou 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: diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 1733d1ff3..fff84ac45 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -3,13 +3,8 @@ from __future__ import annotations 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, - Literal, - Union, -) +from typing import Any from uuid import UUID from langchain_core.runnables.config import RunnableConfig @@ -17,7 +12,7 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text 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, @@ -66,82 +61,43 @@ class CheckpointPayload(TypedDict): tasks: list[CheckpointTask] -class DebugOutputBase(TypedDict): - timestamp: str - step: int - - -class DebugOutputTask(DebugOutputBase): - type: Literal["task"] - payload: TaskPayload - - -class DebugOutputTaskResult(DebugOutputBase): - type: Literal["task_result"] - payload: TaskResultPayload - - -class DebugOutputCheckpoint(DebugOutputBase): - type: Literal["checkpoint"] - payload: CheckpointPayload - - -DebugOutput = Union[DebugOutputTask, DebugOutputTaskResult, DebugOutputCheckpoint] - - 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[TaskPayload]: """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]: +) -> Iterator[TaskResultPayload]: """Produce "task_result" events for stream_mode=debug.""" stream_channels_list = ( [stream_keys] if isinstance(stream_keys, str) else stream_keys ) 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,17 +115,15 @@ 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, output_keys: str | Sequence[str], -) -> Iterator[DebugOutputCheckpoint]: +) -> Iterator[CheckpointPayload]: """Produce "checkpoint" events for stream_mode=debug.""" parent_ns = config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -193,42 +147,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) + ], } diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 9043d6041..07d4a97ad 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -11,6 +11,7 @@ from contextlib import ( AsyncExitStack, ExitStack, ) +from datetime import datetime, timezone from inspect import signature from types import TracebackType from typing import ( @@ -423,7 +424,7 @@ class PregelLoop: ), ): # 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 @@ class PregelLoop: self._emit( "checkpoints", map_debug_checkpoint, - self.step - 1, # printing checkpoint for previous step { **self.checkpoint_config, CONF: { @@ -486,7 +486,6 @@ class PregelLoop: 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 @@ class PregelLoop: 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 @@ class PregelLoop: ) -> 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 @@ class PregelLoop: self._emit( "tasks", map_debug_task_results, - self.step, (task, writes), self.stream_keys, ) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index a5004ab2d..292deee03 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -46,7 +46,9 @@ Checkpointer = Union[None, bool, BaseCheckpointSaver] - 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 @@ StreamMode = Literal["values", "updates", "checkpoints", "tasks", "messages", "c - `"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] diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index c3e135603..5aa413be9 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -36,7 +36,15 @@ Represents the status of a thread: """ 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: @@ -44,6 +52,8 @@ Defines the mode of streaming: - "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. """