Introduce "tasks" and "checkpoints" stream modes (#5117)

This commit is contained in:
Nuno Campos
2025-06-16 10:14:18 -07:00
committed by GitHub
5 changed files with 98 additions and 115 deletions
+2 -1
View File
@@ -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)`.
+48 -101
View File
@@ -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)
],
}
+31 -10
View File
@@ -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,
)
+6 -2
View File
@@ -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]
+11 -1
View File
@@ -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.
"""