From c9dcd3fbb53e36cccf17172609e3c28df99f0dc2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 12:29:22 -0700 Subject: [PATCH] Map debug=True to print_mode=['values', 'updates'] --- libs/langgraph/langgraph/pregel/__init__.py | 66 +++++++++++---------- libs/langgraph/langgraph/pregel/debug.py | 51 +--------------- libs/langgraph/langgraph/pregel/loop.py | 47 +-------------- 3 files changed, 38 insertions(+), 126 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index eb06ce347..1a6823d83 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2203,14 +2203,12 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self, config: RunnableConfig, *, - stream_mode: StreamMode | list[StreamMode], + stream_mode: StreamMode | Sequence[StreamMode], print_mode: StreamMode | Sequence[StreamMode], output_keys: str | Sequence[str] | None, interrupt_before: All | Sequence[str] | None, interrupt_after: All | Sequence[str] | None, - debug: bool | None, ) -> tuple[ - bool, set[StreamMode], str | Sequence[str], All | Sequence[str], @@ -2221,7 +2219,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ]: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") - debug = debug if debug is not None else self.debug if output_keys is None: output_keys = self.stream_channels_asis else: @@ -2258,7 +2255,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou else: cache = self.cache return ( - debug, stream_modes, output_keys, interrupt_before, @@ -2273,7 +2269,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou input: InputT, config: RunnableConfig | None = None, *, - stream_mode: StreamMode | list[StreamMode] | None = None, + stream_mode: StreamMode | Sequence[StreamMode] | None = None, print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, @@ -2309,7 +2305,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. checkpoint_during: Whether to checkpoint intermediate steps, defaults to False. If False, only the final checkpoint is saved. - debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream events from inside subgraphs, defaults to False. If True, the events will be emitted as tuples `(namespace, data)`, or `(namespace, mode, data)` if `stream_mode` is a list, @@ -2330,6 +2325,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou if config is not None and CONFIG_KEY_TASK_ID in config.get(CONF, {}) else self.stream_mode ) + if debug or self.debug: + print_mode = ["updates", "values"] stream = SyncQueue() @@ -2344,7 +2341,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou try: # assign defaults ( - debug, stream_modes, output_keys, interrupt_before_, @@ -2359,7 +2355,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, - debug=debug, ) # set up subgraph checkpointing if self.checkpointer is True: @@ -2407,7 +2402,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou interrupt_before=interrupt_before_, interrupt_after=interrupt_after_, manager=run_manager, - debug=debug, checkpoint_during=checkpoint_during if checkpoint_during is not None else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), @@ -2497,7 +2491,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou input: InputT, config: RunnableConfig | None = None, *, - stream_mode: StreamMode | list[StreamMode] | None = None, + stream_mode: StreamMode | Sequence[StreamMode] | None = None, print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, @@ -2532,7 +2526,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph. interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph. checkpoint_during: Whether to checkpoint intermediate steps, defaults to False. If False, only the final checkpoint is saved. - debug: Whether to print debug information during execution, defaults to False. subgraphs: Whether to stream events from inside subgraphs, defaults to False. If True, the events will be emitted as tuples `(namespace, data)`, or `(namespace, mode, data)` if `stream_mode` is a list, @@ -2553,6 +2546,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou if config is not None and CONFIG_KEY_TASK_ID in config.get(CONF, {}) else self.stream_mode ) + if debug or self.debug: + print_mode = ["updates", "values"] stream = AsyncQueue() aioloop = asyncio.get_running_loop() @@ -2586,7 +2581,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou try: # assign defaults ( - debug, stream_modes, output_keys, interrupt_before_, @@ -2601,7 +2595,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, - debug=debug, ) # set up subgraph checkpointing if self.checkpointer is True: @@ -2652,7 +2645,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou interrupt_before=interrupt_before_, interrupt_after=interrupt_after_, manager=run_manager, - debug=debug, checkpoint_during=checkpoint_during if checkpoint_during is not None else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True), @@ -2748,7 +2740,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - debug: bool | None = None, **kwargs: Any, ) -> dict[str, Any] | Any: """Run the graph with a single input and config. @@ -2761,7 +2752,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou output_keys: Optional. The output keys to retrieve from the graph run. interrupt_before: Optional. The nodes to interrupt the graph run before. interrupt_after: Optional. The nodes to interrupt the graph run after. - debug: Optional. Enable debug mode for the graph run. **kwargs: Additional keyword arguments to pass to the graph run. Returns: @@ -2777,22 +2767,30 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou for chunk in self.stream( input, config, - stream_mode=stream_mode, + stream_mode=["updates", "values"] + if stream_mode == "values" + else stream_mode, print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, - debug=debug, **kwargs, ): if stream_mode == "values": + if len(chunk) == 2: + mode, payload = cast(tuple[StreamMode, Any], chunk) + else: + _, mode, payload = cast( + tuple[tuple[str, ...], StreamMode, Any], chunk + ) if ( - isinstance(chunk, dict) - and (ints := chunk.get(INTERRUPT)) is not None + mode == "updates" + and isinstance(payload, dict) + and (ints := payload.get(INTERRUPT)) is not None ): interrupts.extend(ints) - else: - latest = chunk + elif mode == "values": + latest = payload else: chunks.append(chunk) @@ -2817,7 +2815,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, - debug: bool | None = None, **kwargs: Any, ) -> dict[str, Any] | Any: """Asynchronously invoke the graph on a single input. @@ -2830,7 +2827,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou output_keys: Optional. The output keys to include in the result. Default is None. interrupt_before: Optional. The nodes to interrupt before. Default is None. interrupt_after: Optional. The nodes to interrupt after. Default is None. - debug: Optional. Whether to enable debug mode. Default is None. **kwargs: Additional keyword arguments. Returns: @@ -2847,22 +2843,30 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou async for chunk in self.astream( input, config, - stream_mode=stream_mode, + stream_mode=["updates", "values"] + if stream_mode == "values" + else stream_mode, print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, - debug=debug, **kwargs, ): if stream_mode == "values": + if len(chunk) == 2: + mode, payload = cast(tuple[StreamMode, Any], chunk) + else: + _, mode, payload = cast( + tuple[tuple[str, ...], StreamMode, Any], chunk + ) if ( - isinstance(chunk, dict) - and (ints := chunk.get(INTERRUPT)) is not None + mode == "updates" + and isinstance(payload, dict) + and (ints := payload.get(INTERRUPT)) is not None ): interrupts.extend(ints) - else: - latest = chunk + elif mode == "values": + latest = payload else: chunks.append(chunk) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index d6471c53c..d8d3bbdae 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -1,13 +1,10 @@ from __future__ import annotations -from collections import defaultdict from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import asdict -from pprint import pformat from typing import Any from uuid import UUID -from langchain_core.runnables.config import RunnableConfig from typing_extensions import TypedDict from langgraph.channels.base import BaseChannel @@ -25,7 +22,7 @@ from langgraph.constants import ( ) from langgraph.pregel.io import read_channels from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot -from langgraph.utils.config import patch_checkpoint_map +from langgraph.utils.config import RunnableConfig, patch_checkpoint_map class TaskPayload(TypedDict): @@ -178,52 +175,6 @@ def map_debug_checkpoint( } -def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: - n_tasks = len(next_tasks) - print( - f"{get_colored_text(f'[{step}:tasks]', color='blue')} " - + get_bolded_text( - f"Starting {n_tasks} task{'s' if n_tasks != 1 else ''} for step {step}:\n" - ) - + "\n".join( - f"- {get_colored_text(task.name, 'green')} -> {pformat(task.input)}" - for task in next_tasks - ) - ) - - -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(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( - metadata: CheckpointMetadata, - channels: Mapping[str, BaseChannel], - whitelist: Sequence[str], -) -> None: - step = metadata["step"] - print( - f"{get_colored_text(f'[{step}:checkpoint]', color='blue')} " - + get_bolded_text(f"State at the end of step {step}:\n") - + pformat(read_channels(channels, whitelist), depth=3) - ) - - def tasks_w_writes( tasks: Iterable[PregelTask | PregelExecutableTask], pending_writes: list[PendingWrite] | None, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 9391fb3a9..e5d4e1534 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -91,9 +91,6 @@ from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results, map_debug_tasks, - print_step_checkpoint, - print_step_tasks, - print_step_writes, ) from langgraph.pregel.executor import ( AsyncBackgroundExecutor, @@ -160,7 +157,6 @@ class PregelLoop: interrupt_after: All | Sequence[str] interrupt_before: All | Sequence[str] checkpoint_during: bool - debug: bool retry_policy: Sequence[RetryPolicy] cache_policy: CachePolicy | None @@ -225,7 +221,6 @@ class PregelLoop: interrupt_after: All | Sequence[str] = EMPTY_SEQ, interrupt_before: All | Sequence[str] = EMPTY_SEQ, manager: None | AsyncParentRunManager | ParentRunManager = None, - debug: bool = False, migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, @@ -254,7 +249,6 @@ class PregelLoop: self.retry_policy = retry_policy self.cache_policy = cache_policy self.checkpoint_during = checkpoint_during - self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD) @@ -435,9 +429,6 @@ class PregelLoop: ): # produce debug output self._emit("tasks", map_debug_tasks, [pushed]) - # debug flag - if self.debug: - print_step_tasks(self.step, [pushed]) # save the new task self.tasks[pushed.id] = pushed # match any pending writes to the new task @@ -521,10 +512,6 @@ class PregelLoop: # produce debug output self._emit("tasks", map_debug_tasks, self.tasks.values()) - # debug flag - if self.debug: - print_step_tasks(self.step, list(self.tasks.values())) - # print output for any tasks we applied previous writes to for task in self.tasks.values(): if task.writes: @@ -535,17 +522,6 @@ class PregelLoop: def after_tick(self) -> None: # finish superstep writes = [w for t in self.tasks.values() for w in t.writes] - # debug flag - if self.debug: - print_step_writes( - self.step, - writes, - ( - [self.stream_keys] - if isinstance(self.stream_keys, str) - else self.stream_keys - ), - ) # all tasks have finished self.updated_channels = apply_writes( self.checkpoint, @@ -708,17 +684,6 @@ class PregelLoop: metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) self.checkpoint_metadata = metadata - # debug flag - if self.debug: - print_step_checkpoint( - metadata, - self.channels, - ( - [self.stream_keys] - if isinstance(self.stream_keys, str) - else self.stream_keys - ), - ) # do checkpoint? do_checkpoint = self._checkpointer_put_after_previous is not None and ( exiting or self.checkpoint_during @@ -883,7 +848,7 @@ class PregelLoop: return if writes[0][0] == INTERRUPT: # in loop.py we append a bool to the PUSH task paths to indicate - # whether or not a call was present (that was popped). If so, + # whether or not a call was present. If so, # we don't emit the interrupt as it'll be emitted by the parent if task.path[0] == PUSH and task.path[-1] is True: return @@ -897,11 +862,7 @@ class PregelLoop: ) } ] - stream_modes = self.stream.modes if self.stream else [] - if "updates" in stream_modes: - self._emit("updates", lambda: iter(interrupts)) - elif "values" in stream_modes: - self._emit("values", lambda: iter(interrupts)) + self._emit("updates", lambda: iter(interrupts)) elif writes[0][0] != ERROR: self._emit( "updates", @@ -938,7 +899,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): input_keys: str | Sequence[str] = EMPTY_SEQ, output_keys: str | Sequence[str] = EMPTY_SEQ, stream_keys: str | Sequence[str] = EMPTY_SEQ, - debug: bool = False, migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, @@ -959,7 +919,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): interrupt_after=interrupt_after, interrupt_before=interrupt_before, manager=manager, - debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, retry_policy=retry_policy, @@ -1111,7 +1070,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): input_keys: str | Sequence[str] = EMPTY_SEQ, output_keys: str | Sequence[str] = EMPTY_SEQ, stream_keys: str | Sequence[str] = EMPTY_SEQ, - debug: bool = False, migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, @@ -1132,7 +1090,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): interrupt_after=interrupt_after, interrupt_before=interrupt_before, manager=manager, - debug=debug, migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, retry_policy=retry_policy,