diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5177a9e15..e02de7172 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -68,6 +68,9 @@ from langgraph.constants import ( TAG_HIDDEN, ) from langgraph.pregel.debug import ( + map_debug_checkpoint, + map_debug_task_results, + map_debug_tasks, print_step_checkpoint, print_step_tasks, print_step_writes, @@ -176,7 +179,7 @@ class Channel: ) -StreamMode = Literal["values", "updates"] +StreamMode = Literal["values", "updates", "debug"] class Pregel( @@ -304,12 +307,15 @@ class Pregel( @property def stream_channels_list(self) -> Sequence[str]: + stream_channels = self.stream_channels_asis return ( - [self.stream_channels] - if isinstance(self.stream_channels, str) - else self.stream_channels or [k for k in self.channels] + [stream_channels] if isinstance(stream_channels, str) else stream_channels ) + @property + def stream_channels_asis(self) -> Union[str, Sequence[str]]: + return self.stream_channels or [k for k in self.channels] + def get_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: @@ -322,11 +328,8 @@ class Pregel( _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, for_execution=False ) - values = read_channels(channels, self.stream_channels_list) return StateSnapshot( - values.get(self.stream_channels, None) - if isinstance(self.stream_channels, str) - else values, + read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, ) @@ -343,11 +346,8 @@ class Pregel( _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, for_execution=False ) - values = read_channels(channels, self.stream_channels_list) return StateSnapshot( - values.get(self.stream_channels, None) - if isinstance(self.stream_channels, str) - else values, + read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, ) @@ -362,11 +362,8 @@ class Pregel( _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, for_execution=False ) - values = read_channels(channels, self.stream_channels_list) yield StateSnapshot( - values.get(self.stream_channels, None) - if isinstance(self.stream_channels, str) - else values, + read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, parent_config, @@ -384,11 +381,8 @@ class Pregel( _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, for_execution=False ) - values = read_channels(channels, self.stream_channels_list) yield StateSnapshot( - values.get(self.stream_channels, None) - if isinstance(self.stream_channels, str) - else values, + read_channels(channels, self.stream_channels_asis), tuple(name for name, _ in next_tasks), config, parent_config, @@ -436,6 +430,8 @@ class Pregel( values, RunnableSequence(*writers) if len(writers) > 1 else writers[0], deque(), + None, + [INTERRUPT], ) # execute task task.proc.invoke( @@ -496,6 +492,8 @@ class Pregel( values, RunnableSequence(*writers) if len(writers) > 1 else writers[0], deque(), + None, + [INTERRUPT], ) # execute task await task.proc.ainvoke( @@ -538,11 +536,7 @@ class Pregel( ]: debug = debug if debug is not None else self.debug if output_keys is None: - output_keys = ( - [chan for chan in self.channels] - if self.stream_channels is None - else self.stream_channels - ) + output_keys = self.stream_channels_asis else: validate_keys(output_keys, self.channels) if input_keys is None: @@ -673,6 +667,9 @@ class Pregel( if debug: print_step_tasks(step, next_tasks) + if stream_mode == "debug": + for chunk in map_debug_tasks(step, next_tasks): + yield chunk # prepare tasks with config tasks_w_config = [ @@ -692,7 +689,7 @@ class Pregel( }, ), ) - for name, input, proc, writes, proc_config in next_tasks + for name, input, proc, writes, proc_config, _ in next_tasks ] futures = [ @@ -713,7 +710,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _ in next_tasks: + for _, _, _, writes, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -732,6 +729,10 @@ class Pregel( yield from map_output_values( output_keys, pending_writes, channels ) + elif stream_mode == "debug": + yield from map_debug_task_results( + step, next_tasks, self.stream_channels_list + ) else: yield from map_output_updates(output_keys, next_tasks) @@ -743,6 +744,17 @@ class Pregel( checkpoint_config = self.checkpointer.put( checkpoint_config, checkpoint ) + if stream_mode == "debug": + yield map_debug_checkpoint( + step, + checkpoint_config, + channels, + self.stream_channels_asis, + ) + elif stream_mode == "debug": + yield map_debug_checkpoint( + step, None, channels, self.stream_channels_asis + ) # after execution, check if we should interrupt if _should_interrupt( @@ -762,7 +774,20 @@ class Pregel( and self.checkpointer.at == CheckpointAt.END_OF_RUN ): checkpoint = create_checkpoint(checkpoint, channels) - self.checkpointer.put(checkpoint_config, checkpoint) + checkpoint_config = self.checkpointer.put( + checkpoint_config, checkpoint + ) + if stream_mode == "debug": + yield map_debug_checkpoint( + step, + checkpoint_config, + channels, + self.stream_channels_asis, + ) + elif self.checkpointer is None and stream_mode == "debug": + yield map_debug_checkpoint( + step, None, channels, self.stream_channels_asis + ) except BaseException as e: run_manager.on_chain_error(e) raise @@ -891,6 +916,9 @@ class Pregel( if debug: print_step_tasks(step, next_tasks) + if stream_mode == "debug": + for chunk in map_debug_tasks(step, next_tasks): + yield chunk # prepare tasks with config tasks_w_config = [ @@ -910,7 +938,7 @@ class Pregel( }, ), ) - for name, input, proc, writes, proc_config in next_tasks + for name, input, proc, writes, proc_config, _ in next_tasks ] futures = ( @@ -938,7 +966,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _ in next_tasks: + for _, _, _, writes, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -958,6 +986,11 @@ class Pregel( output_keys, pending_writes, channels ): yield chunk + elif stream_mode == "debug": + for chunk in map_debug_task_results( + step, next_tasks, self.stream_channels_list + ): + yield chunk else: for chunk in map_output_updates(output_keys, next_tasks): yield chunk @@ -970,6 +1003,17 @@ class Pregel( checkpoint_config = await self.checkpointer.aput( checkpoint_config, checkpoint ) + if stream_mode == "debug": + yield map_debug_checkpoint( + step, + checkpoint_config, + channels, + self.stream_channels_asis, + ) + elif stream_mode == "debug": + yield map_debug_checkpoint( + step, None, channels, self.stream_channels_asis + ) # after execution, check if we should interrupt if _should_interrupt( @@ -989,7 +1033,17 @@ class Pregel( and self.checkpointer.at == CheckpointAt.END_OF_RUN ): checkpoint = create_checkpoint(checkpoint, channels) - await self.checkpointer.aput(checkpoint_config, checkpoint) + checkpoint_config = await self.checkpointer.aput( + checkpoint_config, checkpoint + ) + if stream_mode == "debug": + yield map_debug_checkpoint( + step, checkpoint_config, channels, self.stream_channels_asis + ) + elif self.checkpointer is None and stream_mode == "debug": + yield map_debug_checkpoint( + step, None, channels, self.stream_channels_asis + ) except BaseException as e: await run_manager.on_chain_error(e) raise @@ -1154,7 +1208,7 @@ def _should_interrupt( # and any channel written to is in interrupt_nodes list and any( node - for node, _, _, _, config in tasks + for node, _, _, _, config, _ in tasks if ( (not config or TAG_HIDDEN not in config.get("tags")) if interrupt_nodes == "*" @@ -1250,13 +1304,14 @@ def _prepare_next_tasks( for name, proc in processes.items(): seen = checkpoint["versions_seen"][name] # If any of the channels read by this process were updated - if any( - checkpoint["channel_versions"][chan] > seen[chan] + if triggers := [ + chan for chan in proc.triggers if not isinstance( read_channel(channels, chan, return_exception=True), EmptyChannelError ) - ): + and checkpoint["channel_versions"][chan] > seen[chan] + ]: # If all trigger channels subscribed by this process are not empty # then invoke the process with the values of all non-empty channels if isinstance(proc.channels, dict): @@ -1297,7 +1352,9 @@ def _prepare_next_tasks( if for_execution: if node := proc.get_node(): tasks.append( - PregelExecutableTask(name, val, node, deque(), proc.config) + PregelExecutableTask( + name, val, node, deque(), proc.config, triggers + ) ) else: tasks.append(PregelTaskDescription(name, val)) diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index 0857b3074..5fd548344 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -1,13 +1,124 @@ +import json from collections import defaultdict +from datetime import datetime, timezone from pprint import pformat -from typing import Any, Iterator, Mapping, Sequence +from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypedDict, Union +from uuid import UUID, uuid5 +from langchain_core.runnables.config import RunnableConfig from langchain_core.utils.input import get_bolded_text, get_colored_text -from langgraph.channels.base import BaseChannel, EmptyChannelError +from langgraph.channels.base import BaseChannel +from langgraph.constants import TAG_HIDDEN +from langgraph.pregel.io import read_channels from langgraph.pregel.types import PregelExecutableTask +class TaskPayload(TypedDict): + id: str + name: str + input: Any + triggers: list[str] + + +class TaskResultPayload(TypedDict): + id: str + result: list[tuple[str, Any]] + + +class CheckpointPayload(TypedDict): + config: Optional[RunnableConfig] + values: dict[str, Any] + + +class DebugOutputBase(TypedDict): + timestamp: str + step: int + type: str + payload: dict[str, Any] + + +class DebugOutputTask(DebugOutputBase): + type: Literal["task"] + payload: TaskPayload + + +class DebugOutputTaskResult(DebugOutputBase): + type: Literal["task_result"] + payload: TaskResultPayload + + +class DebugOutputCheckpoint(DebugOutputBase): + type: Literal["checkpoint"] + payload: CheckpointPayload + + +DebugOutput = Union[DebugOutputTask, DebugOutputTaskResult, DebugOutputCheckpoint] + + +TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8") + + +def map_debug_tasks( + step: int, tasks: list[PregelExecutableTask] +) -> Iterator[DebugOutputTask]: + ts = datetime.now(timezone.utc).isoformat() + for name, input, _, _, config, triggers in tasks: + if config is not None and TAG_HIDDEN in config.get("tags", []): + continue + + yield { + "type": "task", + "timestamp": ts, + "step": step, + "payload": { + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step)))), + "name": name, + "input": input, + "triggers": triggers, + }, + } + + +def map_debug_task_results( + step: int, + tasks: list[PregelExecutableTask], + stream_channels_list: Sequence[str], +) -> Iterator[DebugOutputTaskResult]: + ts = datetime.now(timezone.utc).isoformat() + for name, _, _, writes, config, _ in tasks: + if config is not None and TAG_HIDDEN in config.get("tags", []): + continue + + yield { + "type": "task_result", + "timestamp": ts, + "step": step, + "payload": { + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step)))), + "result": [w for w in writes if w[0] in stream_channels_list], + }, + } + + +def map_debug_checkpoint( + step: int, + config: RunnableConfig, + channels: Mapping[str, BaseChannel], + stream_channels: Union[str, Sequence[str]], +) -> DebugOutputCheckpoint: + ts = datetime.now(timezone.utc).isoformat() + return { + "type": "checkpoint", + "timestamp": ts, + "step": step, + "payload": { + "config": config, + "values": read_channels(channels, stream_channels), + }, + } + + def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: n_tasks = len(next_tasks) print( @@ -17,7 +128,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: ) + "\n".join( f"- {get_colored_text(name, 'green')} -> {pformat(val)}" - for name, val, _, _, _ in next_tasks + for name, val, _, _, _, _ in next_tasks ) ) @@ -47,16 +158,5 @@ def print_step_checkpoint( 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, - ) + + pformat(read_channels(channels, whitelist), depth=3) ) - - -def _read_channels(channels: Mapping[str, BaseChannel]) -> Iterator[tuple[str, Any]]: - for name, channel in channels.items(): - try: - yield (name, channel.get()) - except EmptyChannelError: - pass diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 7dfc738e5..67382fc77 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -105,7 +105,7 @@ def map_output_updates( if updated := AddableUpdatesDict( { node: value - for node, _, _, writes, _ in output_tasks + for node, _, _, writes, _, _ in output_tasks for chan, value in writes if chan == output_channels } @@ -115,7 +115,7 @@ def map_output_updates( if updated := AddableUpdatesDict( { node: {chan: value for chan, value in writes if chan in output_channels} - for node, _, _, writes, _ in output_tasks + for node, _, _, writes, _, _ in output_tasks if any(chan in output_channels for chan, _ in writes) } ): diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index fd728a7f1..7c6075c98 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -14,7 +14,8 @@ class PregelExecutableTask(NamedTuple): input: Any proc: Runnable writes: deque[tuple[str, Any]] - config: Optional[RunnableConfig] = None + config: Optional[RunnableConfig] + triggers: list[str] class StateSnapshot(NamedTuple): diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 0a4e4cea5..ad4959657 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -15,6 +15,10 @@ def validate_graph( interrupt_after_nodes: Union[All, Sequence[str]], interrupt_before_nodes: Union[All, Sequence[str]], ) -> None: + for chan in channels: + if chan == INTERRUPT: + raise ValueError(f"Channel name {INTERRUPT} is reserved") + subscribed_channels = set[str]() for name, node in nodes.items(): if name == INTERRUPT: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 84d5db4d5..5aac3541b 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -399,6 +399,86 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: {"inbox": [3], "output": 13}, {"inbox": [], "output": 4}, ] + assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="debug")] == [ + { + "type": "task", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa", + "name": "one", + "input": 2, + "triggers": ["input"], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "34e90af0-f97e-54e0-a159-691da37f175f", + "name": "two", + "input": [12], + "triggers": ["inbox"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa", + "result": [("inbox", 3)], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "34e90af0-f97e-54e0-a159-691da37f175f", + "result": [("output", 13)], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": {"config": None, "values": {"output": 13, "inbox": [3]}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "cf7cf374-2a2a-556f-8561-91737af89d2f", + "name": "two", + "input": [3], + "triggers": ["inbox"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "cf7cf374-2a2a-556f-8561-91737af89d2f", + "result": [("output", 4)], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": {"config": None, "values": {"output": 4, "inbox": []}}, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": {"config": None, "values": {"output": 4, "inbox": []}}, + }, + ] def test_batch_two_processes_in_out() -> None: @@ -3413,6 +3493,248 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - with SqliteSaver.from_conn_string(":memory:") as saver: saver.at = checkpoint_at + + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=saver) + thread10 = {"configurable": {"thread_id": "10"}} + if checkpoint_at is CheckpointAt.END_OF_RUN: + assert [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": None, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": None, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": None, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": None, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "config": { + "configurable": { + "thread_id": "10", + "thread_ts": AnyStr(), + } + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + }, + }, + ] + else: + assert [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + }, + }, + ] + tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] ) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index c5e566059..ff099b3e2 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -379,6 +379,88 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: {"inbox": [3], "output": 13}, {"inbox": [], "output": 4}, ] + assert [ + c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="debug") + ] == [ + { + "type": "task", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa", + "name": "one", + "input": 2, + "triggers": ["input"], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "34e90af0-f97e-54e0-a159-691da37f175f", + "name": "two", + "input": [12], + "triggers": ["inbox"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "7a3cc398-2e02-5023-ad7b-e4848d3b67fa", + "result": [("inbox", 3)], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "id": "34e90af0-f97e-54e0-a159-691da37f175f", + "result": [("output", 13)], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": {"config": None, "values": {"output": 13, "inbox": [3]}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "cf7cf374-2a2a-556f-8561-91737af89d2f", + "name": "two", + "input": [3], + "triggers": ["inbox"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "cf7cf374-2a2a-556f-8561-91737af89d2f", + "result": [("output", 4)], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": {"config": None, "values": {"output": 4, "inbox": []}}, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": {"config": None, "values": {"output": 4, "inbox": []}}, + }, + ] async def test_batch_two_processes_in_out() -> None: @@ -2945,6 +3027,250 @@ async def test_branch_then( async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: saver.at = checkpoint_at + + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=saver) + thread10 = {"configurable": {"thread_id": "10"}} + if checkpoint_at is CheckpointAt.END_OF_RUN: + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": None, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": None, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": None, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": None, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "config": { + "configurable": { + "thread_id": "10", + "thread_ts": AnyStr(), + } + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + }, + }, + ] + else: + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + }, + }, + ] + tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] )