From d3353427ae56fd317860ee19c937621566d712e6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 24 Apr 2024 14:55:12 -0700 Subject: [PATCH 1/2] Improve debug logging --- langgraph/pregel/__init__.py | 24 +++++++++++++++---- langgraph/pregel/debug.py | 46 +++++++++++++++++++++++++++++------- langgraph/pregel/log.py | 2 +- 3 files changed, 57 insertions(+), 15 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 0d18bfb2e..117d95fda 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, @@ -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": @@ -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": 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") From c1feff22117d30c615d42c6042af54b7b59c7baf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 24 Apr 2024 15:04:47 -0700 Subject: [PATCH 2/2] Rename for clarity --- langgraph/pregel/__init__.py | 52 ++++++++++++++++++------------------ 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 117d95fda..715d20299 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -521,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, @@ -545,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 @@ -556,8 +556,8 @@ class Pregel( stream_mode, input_keys, output_keys, - interrupt_before_nodes, - interrupt_after_nodes, + interrupt_before, + interrupt_after, ) def stream( @@ -568,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.""" @@ -594,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 @@ -659,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, ): @@ -743,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, ): @@ -778,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) @@ -812,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 @@ -877,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, ): @@ -970,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, ): @@ -1038,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, ): @@ -1094,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, ):