diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 0d18bfb2e..715d20299 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -67,7 +67,11 @@ from langgraph.constants import ( CONFIG_KEY_SEND, INTERRUPT, ) -from langgraph.pregel.debug import print_checkpoint, print_step_start +from langgraph.pregel.debug import ( + print_step_checkpoint, + print_step_tasks, + print_step_writes, +) from langgraph.pregel.io import ( map_input, map_output_updates, @@ -517,8 +521,8 @@ class Pregel( stream_mode: Optional[StreamMode] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before_nodes: Optional[Sequence[str]] = None, - interrupt_after_nodes: Optional[Sequence[str]] = None, + interrupt_before: Optional[Sequence[str]] = None, + interrupt_after: Optional[Sequence[str]] = None, debug: Optional[bool] = None, ) -> tuple[ bool, @@ -541,8 +545,8 @@ class Pregel( input_keys = self.input_channels else: validate_keys(input_keys, self.channels) - interrupt_before_nodes = interrupt_before_nodes or self.interrupt_before_nodes - interrupt_after_nodes = interrupt_after_nodes or self.interrupt_after_nodes + interrupt_before = interrupt_before or self.interrupt_before_nodes + interrupt_after = interrupt_after or self.interrupt_after_nodes stream_mode = stream_mode if stream_mode is not None else self.stream_mode if config is not None and config.get("configurable", {}).get(CONFIG_KEY_READ): # if being called as a node in another graph, always use values mode @@ -552,8 +556,8 @@ class Pregel( stream_mode, input_keys, output_keys, - interrupt_before_nodes, - interrupt_after_nodes, + interrupt_before, + interrupt_after, ) def stream( @@ -564,8 +568,8 @@ class Pregel( stream_mode: Optional[StreamMode] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before_nodes: Optional[Sequence[str]] = None, - interrupt_after_nodes: Optional[Sequence[str]] = None, + interrupt_before: Optional[Sequence[str]] = None, + interrupt_after: Optional[Sequence[str]] = None, debug: Optional[bool] = None, ) -> Iterator[Union[dict[str, Any], Any]]: """Stream graph steps for a single input.""" @@ -590,15 +594,15 @@ class Pregel( stream_mode, input_keys, output_keys, - interrupt_before_nodes, - interrupt_after_nodes, + interrupt_before, + interrupt_after, ) = self._defaults( config, stream_mode=stream_mode, input_keys=input_keys, output_keys=output_keys, - interrupt_before_nodes=interrupt_before_nodes, - interrupt_after_nodes=interrupt_after_nodes, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, debug=debug, ) # copy nodes to ignore mutations during execution @@ -655,7 +659,7 @@ class Pregel( # before execution, check if we should interrupt if _should_interrupt( checkpoint, - interrupt_before_nodes, + interrupt_before, self.stream_channels_list, next_tasks, ): @@ -664,7 +668,7 @@ class Pregel( checkpoint = next_checkpoint if debug: - print_step_start(step, next_tasks) + print_step_tasks(step, next_tasks) # prepare tasks with config tasks_w_config = [ @@ -708,11 +712,16 @@ class Pregel( for _, _, _, writes, _ in next_tasks: pending_writes.extend(writes) + if debug: + print_step_writes( + step, pending_writes, self.stream_channels_list + ) + # apply writes to channels _apply_writes(checkpoint, channels, pending_writes) if debug: - print_checkpoint(step, channels) + print_step_checkpoint(step, channels, self.stream_channels_list) # yield current value or updates if stream_mode == "values": @@ -734,7 +743,7 @@ class Pregel( # after execution, check if we should interrupt if _should_interrupt( checkpoint, - interrupt_after_nodes, + interrupt_after, self.stream_channels_list, next_tasks, ): @@ -769,8 +778,8 @@ class Pregel( stream_mode: Optional[StreamMode] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before_nodes: Optional[Sequence[str]] = None, - interrupt_after_nodes: Optional[Sequence[str]] = None, + interrupt_before: Optional[Sequence[str]] = None, + interrupt_after: Optional[Sequence[str]] = None, debug: Optional[bool] = None, ) -> AsyncIterator[Union[dict[str, Any], Any]]: config = ensure_config(config) @@ -803,15 +812,15 @@ class Pregel( stream_mode, input_keys, output_keys, - interrupt_before_nodes, - interrupt_after_nodes, + interrupt_before, + interrupt_after, ) = self._defaults( config, stream_mode=stream_mode, input_keys=input_keys, output_keys=output_keys, - interrupt_before_nodes=interrupt_before_nodes, - interrupt_after_nodes=interrupt_after_nodes, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, debug=debug, ) # copy nodes to ignore mutations during execution @@ -868,7 +877,7 @@ class Pregel( # before execution, check if we should interrupt if _should_interrupt( checkpoint, - interrupt_before_nodes, + interrupt_before, self.stream_channels_list, next_tasks, ): @@ -877,7 +886,7 @@ class Pregel( checkpoint = next_checkpoint if debug: - print_step_start(step, next_tasks) + print_step_tasks(step, next_tasks) # prepare tasks with config tasks_w_config = [ @@ -928,11 +937,16 @@ class Pregel( for _, _, _, writes, _ in next_tasks: pending_writes.extend(writes) + if debug: + print_step_writes( + step, pending_writes, self.stream_channels_list + ) + # apply writes to channels _apply_writes(checkpoint, channels, pending_writes) if debug: - print_checkpoint(step, channels) + print_step_checkpoint(step, channels, self.stream_channels_list) # yield current value or updates if stream_mode == "values": @@ -956,7 +970,7 @@ class Pregel( # after execution, check if we should interrupt if _should_interrupt( checkpoint, - interrupt_after_nodes, + interrupt_after, self.stream_channels_list, next_tasks, ): @@ -1024,8 +1038,8 @@ class Pregel( stream_mode=stream_mode, output_keys=output_keys, input_keys=input_keys, - interrupt_before_nodes=interrupt_before_nodes, - interrupt_after_nodes=interrupt_after_nodes, + interrupt_before=interrupt_before_nodes, + interrupt_after=interrupt_after_nodes, debug=debug, **kwargs, ): @@ -1080,8 +1094,8 @@ class Pregel( stream_mode=stream_mode, output_keys=output_keys, input_keys=input_keys, - interrupt_before_nodes=interrupt_before_nodes, - interrupt_after_nodes=interrupt_after_nodes, + interrupt_before=interrupt_before_nodes, + interrupt_after=interrupt_after_nodes, debug=debug, **kwargs, ): diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index 64a692e52..0857b3074 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -1,5 +1,6 @@ +from collections import defaultdict from pprint import pformat -from typing import Any, Iterator, Mapping +from typing import Any, Iterator, Mapping, Sequence from langchain_core.utils.input import get_bolded_text, get_colored_text @@ -7,22 +8,49 @@ from langgraph.channels.base import BaseChannel, EmptyChannelError from langgraph.pregel.types import PregelExecutableTask -def print_step_start(step: int, next_tasks: list[PregelExecutableTask]) -> None: +def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: n_tasks = len(next_tasks) print( - f"{get_colored_text('[langgraph/step]', color='blue')} " + f"{get_colored_text(f'[{step}:tasks]', color='blue')} " + get_bolded_text( - f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Next tasks:\n" + f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}:\n" + ) + + "\n".join( + f"- {get_colored_text(name, 'green')} -> {pformat(val)}" + for name, val, _, _, _ in next_tasks ) - + "\n".join(f"- {name}({pformat(val)})" for name, val, _, _, _ in next_tasks) ) -def print_checkpoint(step: int, channels: Mapping[str, BaseChannel]) -> None: +def print_step_writes( + step: int, writes: Sequence[tuple[str, Any]], whitelist: Sequence[str] +) -> None: + by_channel: dict[str, list[Any]] = defaultdict(list) + for channel, value in writes: + if channel in whitelist: + by_channel[channel].append(value) print( - f"{get_colored_text('[langgraph/checkpoint]', color='blue')} " - + get_bolded_text(f"Finishing step {step}. Channel values:\n") - + pformat({name: val for name, val in _read_channels(channels)}, depth=3) + f"{get_colored_text(f'[{step}:writes]', color='blue')} " + + get_bolded_text( + f"Finished step {step} with writes to {len(by_channel)} channel{'s' if len(by_channel) > 1 else ''}:\n" + ) + + "\n".join( + f"- {get_colored_text(name, 'yellow')} -> {', '.join(pformat(v) for v in vals)}" + for name, vals in by_channel.items() + ) + ) + + +def print_step_checkpoint( + step: int, channels: Mapping[str, BaseChannel], whitelist: Sequence[str] +) -> None: + print( + f"{get_colored_text(f'[{step}:checkpoint]', color='blue')} " + + get_bolded_text(f"State at the end of step {step}:\n") + + pformat( + {name: val for name, val in _read_channels(channels) if name in whitelist}, + depth=3, + ) ) diff --git a/langgraph/pregel/log.py b/langgraph/pregel/log.py index eea436a37..fd127777b 100644 --- a/langgraph/pregel/log.py +++ b/langgraph/pregel/log.py @@ -1,3 +1,3 @@ import logging -logger = logging.getLogger(__name__) +logger = logging.getLogger("langgraph")