mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-25 09:02:25 +02:00
Introduce "tasks" and "checkpoints" stream modes (#5117)
This commit is contained in:
@@ -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)`.
|
||||
|
||||
@@ -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)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user