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,