diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3ee6420fa..3946ac746 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2298,7 +2298,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/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 ee4382557..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 ( @@ -119,6 +120,7 @@ from langgraph.types import ( PregelScratchpad, RetryPolicy, StreamChunk, + StreamMode, StreamProtocol, ) from langgraph.utils.config import patch_configurable @@ -422,7 +424,7 @@ class PregelLoop: ), ): # produce debug output - self._emit("debug", map_debug_tasks, self.step, [pushed]) + self._emit("tasks", map_debug_tasks, [pushed]) # debug flag if self.debug: print_step_tasks(self.step, [pushed]) @@ -472,9 +474,8 @@ 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 { **self.checkpoint_config, CONF: { @@ -485,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, @@ -509,7 +509,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.tasks.values()) # debug flag if self.debug: @@ -834,17 +834,39 @@ class PregelLoop: def _emit( self, - mode: str, + mode: StreamMode, values: Callable[P, Iterator[Any]], *args: P.args, **kwargs: P.kwargs, ) -> 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 @@ -885,9 +907,8 @@ class PregelLoop: ) if not cached: self._emit( - "debug", + "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 3e907e3eb..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", "debug", "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. @@ -55,7 +57,9 @@ 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. +- `"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. """