Merge pull request #343 from langchain-ai/nc/24apr/dbeug

Improve debug logging
This commit is contained in:
Nuno Campos
2024-04-24 15:36:18 -07:00
committed by GitHub
3 changed files with 83 additions and 41 deletions
+45 -31
View File
@@ -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,
@@ -517,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,
@@ -541,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
@@ -552,8 +556,8 @@ class Pregel(
stream_mode,
input_keys,
output_keys,
interrupt_before_nodes,
interrupt_after_nodes,
interrupt_before,
interrupt_after,
)
def stream(
@@ -564,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."""
@@ -590,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
@@ -655,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,
):
@@ -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":
@@ -734,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,
):
@@ -769,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)
@@ -803,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
@@ -868,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,
):
@@ -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":
@@ -956,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,
):
@@ -1024,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,
):
@@ -1080,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,
):
+37 -9
View File
@@ -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,
)
)
+1 -1
View File
@@ -1,3 +1,3 @@
import logging
logger = logging.getLogger(__name__)
logger = logging.getLogger("langgraph")