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

Merge pull request #370 from langchain-ai/nc/30apr/stream-mode-debug

Browse Source 
Implement stream_mode=debug
This commit is contained in:
Nuno Campos
2024-04-30 16:54:37 -07:00
committed by GitHub
parent 49c95cefcf 5e322a6550
commit 7d3133b239
7 changed files with 864 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+93 -36
View File
@@ -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))
langgraph/pregel/debug.py
+115 -15
View File
@@ -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
langgraph/pregel/io.py
+2 -2
View File
@@ -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)
}
):
langgraph/pregel/types.py
+2 -1
View File
@@ -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):
langgraph/pregel/validate.py
+4
View File
@@ -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:
tests/test_pregel.py
+322
View File
@@ -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"]
)
tests/test_pregel_async.py
+326
View File
@@ -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"]
)