WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add print_mode= arg to invoke/stream

Browse Source 
- This is more flexible version of the debug= flag, which we'll deprecate
This commit is contained in:
Nuno Campos
2025-06-25 11:52:42 -07:00
parent 0aefe68a5f
commit e23da72ccd
4 changed files with 130 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+104 -44
View File
@@ -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
libs/langgraph/langgraph/pregel/debug.py
+23 -1
View File
@@ -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"
libs/langgraph/langgraph/utils/queue.py
+1 -1
View File
@@ -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),
libs/langgraph/tests/test_pregel.py
+2 -1
View File
@@ -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,