From e23da72ccd558f70b85f92c3f34d57b8cf06fc07 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 11:52:42 -0700 Subject: [PATCH 1/2] Add print_mode= arg to invoke/stream - This is more flexible version of the debug= flag, which we'll deprecate --- libs/langgraph/langgraph/pregel/__init__.py | 148 ++++++++++++++------ libs/langgraph/langgraph/pregel/debug.py | 24 +++- libs/langgraph/langgraph/utils/queue.py | 2 +- libs/langgraph/tests/test_pregel.py | 3 +- 4 files changed, 130 insertions(+), 47 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2f9ca1425..eb06ce347 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -82,7 +82,7 @@ from langgraph.pregel.checkpoint import ( create_checkpoint, empty_checkpoint, ) -from langgraph.pregel.debug import tasks_w_writes +from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes from langgraph.pregel.draw import draw_graph from langgraph.pregel.io import map_input, read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop @@ -2203,7 +2203,8 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou self, config: RunnableConfig, *, - stream_mode: StreamMode | list[StreamMode] | None, + stream_mode: StreamMode | list[StreamMode], + print_mode: StreamMode | Sequence[StreamMode], output_keys: str | Sequence[str] | None, interrupt_before: All | Sequence[str] | None, interrupt_after: All | Sequence[str] | None, @@ -2227,14 +2228,14 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou validate_keys(output_keys, self.channels) interrupt_before = interrupt_before or self.interrupt_before_nodes interrupt_after = interrupt_after or self.interrupt_after_nodes - if stream_mode is None and CONFIG_KEY_TASK_ID in config.get(CONF, {}): - # if being called as a node in another graph, default to values mode - # but don't overwrite stream_mode arg if provided - stream_mode = ["values"] - elif stream_mode is None: - stream_mode = self.stream_mode if not isinstance(stream_mode, list): - stream_mode = [stream_mode] + stream_modes = {stream_mode} + else: + stream_modes = set(stream_mode) + if isinstance(print_mode, str): + stream_modes.add(print_mode) + else: + stream_modes.update(print_mode) if self.checkpointer is False: checkpointer: BaseCheckpointSaver | None = None elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}): @@ -2258,7 +2259,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou cache = self.cache return ( debug, - set(stream_mode), + stream_modes, output_keys, interrupt_before, interrupt_after, @@ -2273,6 +2274,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou config: RunnableConfig | None = None, *, stream_mode: StreamMode | list[StreamMode] | None = None, + print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, @@ -2302,6 +2304,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou The streamed outputs will be tuples of `(mode, data)`. See [LangGraph streaming guide](https://langchain-ai.github.io/langgraph/how-tos/streaming/) for more details. + print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. output_keys: The keys to stream, defaults to all non-context channels. 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. @@ -2319,22 +2322,16 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou The output of each step in the graph. The output shape depends on the stream_mode. """ - stream = SyncQueue() + if stream_mode is None: + # if being called as a node in another graph, default to values mode + # but don't overwrite stream_mode arg if provided + stream_mode = ( + "values" + if config is not None and CONFIG_KEY_TASK_ID in config.get(CONF, {}) + else self.stream_mode + ) - def output() -> Iterator: - while True: - try: - ns, mode, payload = stream.get(block=False) - except queue.Empty: - break - if subgraphs and isinstance(stream_mode, list): - yield (ns, mode, payload) - elif isinstance(stream_mode, list): - yield (mode, payload) - elif subgraphs: - yield (ns, payload) - else: - yield payload + stream = SyncQueue() config = ensure_config(self.config, config) callback_manager = get_callback_manager_for_config(config) @@ -2358,6 +2355,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) = self._defaults( config, stream_mode=stream_mode, + print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, @@ -2469,10 +2467,14 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou schedule_task=loop.accept_push, ): # emit output - yield from output() + yield from _output( + stream_mode, print_mode, subgraphs, stream.get, queue.Empty + ) loop.after_tick() # emit output - yield from output() + yield from _output( + stream_mode, print_mode, subgraphs, stream.get, queue.Empty + ) # handle exit if loop.status == "out_of_steps": msg = create_error_message( @@ -2496,6 +2498,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou config: RunnableConfig | None = None, *, stream_mode: StreamMode | list[StreamMode] | None = None, + print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, @@ -2524,6 +2527,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou The streamed outputs will be tuples of `(mode, data)`. See [LangGraph streaming guide](https://langchain-ai.github.io/langgraph/how-tos/streaming/) for more details. + print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. output_keys: The keys to stream, defaults to all non-context channels. 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. @@ -2541,6 +2545,15 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou The output of each step in the graph. The output shape depends on the stream_mode. """ + if stream_mode is None: + # if being called as a node in another graph, default to values mode + # but don't overwrite stream_mode arg if provided + stream_mode = ( + "values" + if config is not None and CONFIG_KEY_TASK_ID in config.get(CONF, {}) + else self.stream_mode + ) + stream = AsyncQueue() aioloop = asyncio.get_running_loop() stream_put = cast( @@ -2548,21 +2561,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou partial(aioloop.call_soon_threadsafe, stream.put_nowait), ) - def output() -> Iterator: - while True: - try: - ns, mode, payload = stream.get_nowait() - except asyncio.QueueEmpty: - break - if subgraphs and isinstance(stream_mode, list): - yield (ns, mode, payload) - elif isinstance(stream_mode, list): - yield (mode, payload) - elif subgraphs: - yield (ns, payload) - else: - yield payload - config = ensure_config(self.config, config) callback_manager = get_async_callback_manager_for_config(config) run_manager = await callback_manager.on_chain_start( @@ -2599,6 +2597,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou ) = self._defaults( config, stream_mode=stream_mode, + print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, @@ -2704,11 +2703,23 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou schedule_task=loop.aaccept_push, ): # emit output - for o in output(): + for o in _output( + stream_mode, + print_mode, + subgraphs, + stream.get_nowait, + asyncio.QueueEmpty, + ): yield o loop.after_tick() # emit output - for o in output(): + for o in _output( + stream_mode, + print_mode, + subgraphs, + stream.get_nowait, + asyncio.QueueEmpty, + ): yield o # handle exit if loop.status == "out_of_steps": @@ -2733,6 +2744,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou config: RunnableConfig | None = None, *, stream_mode: StreamMode = "values", + print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, @@ -2745,6 +2757,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou input: The input data for the graph. It can be a dictionary or any other type. config: Optional. The configuration for the graph run. stream_mode: Optional[str]. The stream mode for the graph run. Default is "values". + print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. 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. @@ -2765,6 +2778,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou input, config, stream_mode=stream_mode, + print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, @@ -2799,6 +2813,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou config: RunnableConfig | None = None, *, stream_mode: StreamMode = "values", + print_mode: StreamMode | Sequence[StreamMode] = (), output_keys: str | Sequence[str] | None = None, interrupt_before: All | Sequence[str] | None = None, interrupt_after: All | Sequence[str] | None = None, @@ -2811,6 +2826,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou input: The input data for the computation. It can be a dictionary or any other type. config: Optional. The configuration for the computation. stream_mode: Optional. The stream mode for the computation. Default is "values". + print_mode: Accepts the same values as `stream_mode`, but only prints the output to the console, for debugging purposes. Does not affect the output of the graph in any way. 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. @@ -2832,6 +2848,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou input, config, stream_mode=stream_mode, + print_mode=print_mode, output_keys=output_keys, interrupt_before=interrupt_before, interrupt_after=interrupt_after, @@ -2906,3 +2923,46 @@ def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str for trigger in node.triggers: trigger_to_nodes[trigger].append(name) return dict(trigger_to_nodes) + + +def _output( + stream_mode: StreamMode | Sequence[StreamMode], + print_mode: StreamMode | Sequence[StreamMode], + stream_subgraphs: bool, + getter: Callable[[], tuple[tuple[str, ...], str, Any]], + empty_exc: type[Exception], +) -> Iterator: + while True: + try: + ns, mode, payload = getter() + except empty_exc: + break + if mode in print_mode: + if stream_subgraphs and ns: + print( + " ".join( + ( + get_bolded_text(f"[{mode}]"), + get_colored_text(f"[graph={ns}]", color="yellow"), + repr(payload), + ) + ) + ) + else: + print( + " ".join( + ( + get_bolded_text(f"[{mode}]"), + repr(payload), + ) + ) + ) + if mode in stream_mode: + if stream_subgraphs and isinstance(stream_mode, list): + yield (ns, mode, payload) + elif isinstance(stream_mode, list): + yield (mode, payload) + elif stream_subgraphs: + yield (ns, payload) + else: + yield payload diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index fff84ac45..d6471c53c 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -8,7 +8,6 @@ from typing import Any from uuid import UUID from langchain_core.runnables.config import RunnableConfig -from langchain_core.utils.input import get_bolded_text, get_colored_text from typing_extensions import TypedDict from langgraph.channels.base import BaseChannel @@ -294,3 +293,26 @@ def tasks_w_writes( ) ) return tuple(out) + + +COLOR_MAPPING = { + "black": "0;30", + "red": "0;31", + "green": "0;32", + "yellow": "0;33", + "blue": "0;34", + "magenta": "0;35", + "cyan": "0;36", + "white": "0;37", + "gray": "1;30", +} + + +def get_colored_text(text: str, color: str) -> str: + """Get colored text.""" + return f"\033[1;3{COLOR_MAPPING[color]}m{text}\033[0m" + + +def get_bolded_text(text: str) -> str: + """Get bolded text.""" + return f"\033[1m{text}\033[0m" diff --git a/libs/langgraph/langgraph/utils/queue.py b/libs/langgraph/langgraph/utils/queue.py index 14ff7875b..c0717fe34 100644 --- a/libs/langgraph/langgraph/utils/queue.py +++ b/libs/langgraph/langgraph/utils/queue.py @@ -91,7 +91,7 @@ class SyncQueue: self._queue.append(item) self._count.release() - def get(self, block=True, timeout=None): + def get(self, block=False, timeout=None): """Remove and return an item from the queue. If optional args 'block' is true and 'timeout' is None (the default), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 47d91bc85..bc32e2480 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3169,7 +3169,8 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.get_graph(xray=True).draw_mermaid() == snapshot assert app.invoke( - {"my_key": "my value", "never_called": never_called}, debug=True + {"my_key": "my value", "never_called": never_called}, + print_mode=["values", "updates"], ) == { "my_key": "my value there and back again", "never_called": never_called, From c9dcd3fbb53e36cccf17172609e3c28df99f0dc2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 25 Jun 2025 12:29:22 -0700 Subject: [PATCH 2/2] 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,