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

Add interrupts and errors to debug stream mode

Browse Source 
- move remaining output code from Pregel.stream to PregelLoop
This commit is contained in:
Nuno Campos
2024-08-21 16:23:48 -07:00
parent 9147d05cc4
commit d85e267a83
6 changed files with 272 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+49 -73
View File
@@ -76,16 +76,12 @@ from langgraph.pregel.algo import (
should_interrupt,
)
from langgraph.pregel.debug import (
map_debug_task_results,
print_step_checkpoint,
print_step_tasks,
print_step_writes,
tasks_w_writes,
)
from langgraph.pregel.io import (
map_output_updates,
read_channels,
)
from langgraph.pregel.io import read_channels
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
@@ -962,6 +958,7 @@ class Pregel(
nodes=self.nodes,
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
) as loop:
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
@@ -970,7 +967,6 @@ class Pregel(
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
manager=run_manager,
@@ -1039,26 +1035,17 @@ class Pregel(
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
# yield updates output for the finished task
if "updates" in stream_modes:
yield from _with_mode(
"updates",
isinstance(stream_mode, list),
map_output_updates(output_keys, [task]),
)
if "debug" in stream_modes:
yield from _with_mode(
"debug",
isinstance(stream_mode, list),
map_debug_task_results(
loop.step,
[task],
self.stream_channels_list,
),
)
else:
# remove references to loop vars
del fut, task
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
if _should_stop_others(done):
break
@@ -1073,21 +1060,21 @@ class Pregel(
[w for t in loop.tasks for w in t.writes],
self.stream_channels_list,
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
run_manager.on_chain_end(loop.output)
except BaseException as e:
@@ -1219,6 +1206,7 @@ class Pregel(
nodes=self.nodes,
specs=self.channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
) as loop:
aioloop = asyncio.get_event_loop()
# Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1228,7 +1216,6 @@ class Pregel(
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
manager=run_manager,
@@ -1298,28 +1285,17 @@ class Pregel(
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
# yield updates output for the finished task
if "updates" in stream_modes:
for chunk in _with_mode(
"updates",
isinstance(stream_mode, list),
map_output_updates(output_keys, [task]),
):
yield chunk
if "debug" in stream_modes:
for chunk in _with_mode(
"debug",
isinstance(stream_mode, list),
map_debug_task_results(
loop.step,
[task],
self.stream_channels_list,
),
):
yield chunk
else:
# remove references to loop vars
del fut, task
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
if _should_stop_others(done):
break
@@ -1334,21 +1310,21 @@ class Pregel(
[w for t in loop.tasks for w in t.writes],
self.stream_channels_list,
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# emit output
while loop.stream:
mode, payload = loop.stream.popleft()
if mode in stream_modes:
if isinstance(stream_mode, list):
yield (mode, payload)
else:
yield payload
# handle exit
if loop.status == "out_of_steps":
raise GraphRecursionError(
f"Recursion limit of {config['recursion_limit']} reached "
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
await run_manager.on_chain_end(loop.output)
except BaseException as e:
libs/langgraph/langgraph/pregel/debug.py
+12 -4
View File
@@ -1,5 +1,6 @@
import json
from collections import defaultdict
from dataclasses import asdict
from datetime import datetime, timezone
from pprint import pformat
from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypedDict, Union
@@ -25,6 +26,8 @@ class TaskPayload(TypedDict):
class TaskResultPayload(TypedDict):
id: str
name: str
error: Optional[str]
interrupts: list[dict]
result: list[tuple[str, Any]]
@@ -97,11 +100,14 @@ def map_debug_tasks(
def map_debug_task_results(
step: int,
tasks: list[PregelExecutableTask],
stream_channels_list: Sequence[str],
tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
stream_keys: Union[str, Sequence[str]],
) -> Iterator[DebugOutputTaskResult]:
stream_channels_list = (
[stream_keys] if isinstance(stream_keys, str) else stream_keys
)
ts = datetime.now(timezone.utc).isoformat()
for name, _, _, writes, config, _, _, _ in tasks:
for (name, _, _, _, config, _, _, _), writes in tasks:
if config is not None and TAG_HIDDEN in config.get("tags", []):
continue
@@ -116,7 +122,9 @@ def map_debug_task_results(
"payload": {
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
"name": name,
"error": next((w[1] for w in writes if w[0] == ERROR), None),
"result": [w for w in writes if w[0] in stream_channels_list],
"interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
},
}
@@ -150,7 +158,7 @@ def map_debug_checkpoint(
else {
"id": t.id,
"name": t.name,
"interrupts": t.interrupts,
"interrupts": tuple(asdict(i) for i in t.interrupts),
}
for t in tasks_w_writes(tasks, pending_writes)
],
libs/langgraph/langgraph/pregel/io.py
+11 -8
View File
@@ -3,7 +3,7 @@ from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union
from langchain_core.runnables.utils import AddableDict
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.constants import TAG_HIDDEN
from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
from langgraph.pregel.log import logger
from langgraph.pregel.types import PregelExecutableTask
@@ -95,19 +95,22 @@ class AddableUpdatesDict(AddableDict):
def map_output_updates(
output_channels: Union[str, Sequence[str]],
tasks: list[PregelExecutableTask],
tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
output_tasks = [
t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
(t, ww)
for t, ww in tasks
if (not t.config or TAG_HIDDEN not in t.config.get("tags"))
and all(k not in (ERROR, INTERRUPT) for k, _ in ww)
]
if not output_tasks:
return
if isinstance(output_channels, str):
updated = [
(task.name, value)
for task in output_tasks
for chan, value in task.writes
for task, writes in output_tasks
for chan, value in writes
if chan == output_channels
]
else:
@@ -116,10 +119,10 @@ def map_output_updates(
task.name,
{chan: value for chan, value in task.writes if chan in output_channels},
)
for task in output_tasks
if any(chan in output_channels for chan, _ in task.writes)
for task, writes in output_tasks
if any(chan in output_channels for chan, _ in writes)
]
grouped = {t.name: [] for t in output_tasks}
grouped = {t.name: [] for t, _ in output_tasks}
for node, value in updated:
grouped[node].append(value)
for node, value in grouped.items():
libs/langgraph/langgraph/pregel/loop.py
+31 -5
View File
@@ -56,7 +56,11 @@ from langgraph.pregel.algo import (
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
from langgraph.pregel.debug import (
map_debug_checkpoint,
map_debug_task_results,
map_debug_tasks,
)
from langgraph.pregel.executor import (
AsyncBackgroundExecutor,
BackgroundExecutor,
@@ -90,6 +94,7 @@ class PregelLoop:
nodes: Mapping[str, PregelNode]
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
output_keys: Union[str, Sequence[str]]
stream_keys: Union[str, Sequence[str]]
is_nested: bool
checkpointer_get_next_version: Callable[[Optional[V]], V]
@@ -138,6 +143,7 @@ class PregelLoop:
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]],
stream_keys: Union[str, Sequence[str]],
) -> None:
self.stream = deque()
self.input = input
@@ -147,6 +153,7 @@ class PregelLoop:
self.nodes = nodes
self.specs = specs
self.output_keys = output_keys
self.stream_keys = stream_keys
self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
@@ -172,12 +179,24 @@ class PregelLoop:
writes,
task_id,
)
if task := next((t for t in self.tasks if t.id == task_id), None):
self.stream.extend(
("updates", v)
for v in map_output_updates(self.output_keys, [(task, writes)])
)
self.stream.extend(
("debug", v)
for v in map_debug_task_results(
self.step, [(task, writes)], self.stream_keys
)
)
else:
raise RuntimeError
def tick(
self,
*,
input_keys: Union[str, Sequence[str]],
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
interrupt_after: Sequence[str] = EMPTY_SEQ,
interrupt_before: Sequence[str] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
@@ -213,7 +232,11 @@ class PregelLoop:
self._put_checkpoint(
{
"source": "loop",
"writes": single(map_output_updates(self.output_keys, self.tasks)),
"writes": single(
map_output_updates(
self.output_keys, [(t, t.writes) for t in self.tasks]
)
),
}
)
# after execution, check if we should interrupt
@@ -256,7 +279,7 @@ class PregelLoop:
self.step - 1, # printing checkpoint for previous step
self.checkpoint_config,
self.channels,
stream_keys,
self.stream_keys,
self.checkpoint_metadata,
self.checkpoint,
self.tasks,
@@ -281,7 +304,6 @@ class PregelLoop:
if all(task.writes for task in self.tasks):
return self.tick(
input_keys=input_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
@@ -435,6 +457,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(
input,
@@ -444,6 +467,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
nodes=nodes,
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
)
self.stack = ExitStack()
if checkpointer:
@@ -522,6 +546,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(
input,
@@ -531,6 +556,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
nodes=nodes,
specs=specs,
output_keys=output_keys,
stream_keys=stream_keys,
)
self.store = AsyncBatchedStore(self.store) if self.store else None
self.stack = AsyncExitStack()
libs/langgraph/tests/test_pregel.py
+21 -2
View File
@@ -769,7 +769,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
]
# forking from any previous checkpoint w/out forking should do nothing
# re-running from any previous checkpoint w/out forking should do nothing
assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == []
assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == []
@@ -1033,6 +1033,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
"name": "one",
"result": [("inbox", 3)],
"error": None,
"interrupts": [],
},
},
{
@@ -1043,6 +1045,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
"name": "two",
"result": [("output", 13)],
"error": None,
"interrupts": [],
},
},
{
@@ -1064,6 +1068,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
"name": "two",
"result": [("output", 4)],
"error": None,
"interrupts": [],
},
},
]
@@ -1744,7 +1750,6 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert cleanup.call_count == 0
for i, chunk in enumerate(app.stream(2)):
assert setup.call_count == 1, "Expected setup to be called once"
assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
@@ -6017,6 +6022,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
"name": "rewrite_query",
"result": [("query", "query: what is weather in sf")],
"error": None,
"interrupts": [],
},
},
),
@@ -6071,6 +6078,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
"name": "retriever_two",
"result": [("docs", ["doc3", "doc4"])],
"error": None,
"interrupts": [],
},
},
),
@@ -6088,6 +6097,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
"name": "retriever_one",
"result": [("docs", ["doc1", "doc2"])],
"error": None,
"interrupts": [],
},
},
),
@@ -6127,6 +6138,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
"name": "qa",
"result": [("answer", "doc1,doc2,doc3,doc4")],
"error": None,
"interrupts": [],
},
},
),
@@ -6544,6 +6557,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
"name": "prepare",
"result": [("my_key", " prepared")],
"error": None,
"interrupts": [],
},
},
{
@@ -6596,6 +6611,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
"error": None,
"interrupts": [],
},
},
{
@@ -6646,6 +6663,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
"name": "finish",
"result": [("my_key", " finished")],
"error": None,
"interrupts": [],
},
},
{
libs/langgraph/tests/test_pregel_async.py
+148 -5
View File
@@ -1278,6 +1278,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
"name": "one",
"result": [("inbox", 3)],
"error": None,
"interrupts": [],
},
},
{
@@ -1288,6 +1290,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
"name": "two",
"result": [("output", 13)],
"error": None,
"interrupts": [],
},
},
{
@@ -1309,6 +1313,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
"name": "two",
"result": [("output", 4)],
"error": None,
"interrupts": [],
},
},
]
@@ -1978,7 +1984,6 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert setup_sync.call_count == 0, "Sync context manager should not be used"
assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
assert setup_async.call_count == 1, "Expected setup to be called once"
assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet"
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
@@ -4744,6 +4749,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
"name": "rewrite_query",
"result": [("query", "query: what is weather in sf")],
"error": None,
"interrupts": [],
},
},
),
@@ -4798,6 +4805,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
"name": "retriever_two",
"result": [("docs", ["doc3", "doc4"])],
"error": None,
"interrupts": [],
},
},
),
@@ -4815,6 +4824,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
"name": "retriever_one",
"result": [("docs", ["doc1", "doc2"])],
"error": None,
"interrupts": [],
},
},
),
@@ -4854,6 +4865,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
"name": "qa",
"result": [("answer", "doc1,doc2,doc3,doc4")],
"error": None,
"interrupts": [],
},
},
),
@@ -5218,6 +5231,8 @@ async def test_branch_then() -> None:
"id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
"name": "prepare",
"result": [("my_key", " prepared")],
"error": None,
"interrupts": [],
},
},
{
@@ -5270,6 +5285,8 @@ async def test_branch_then() -> None:
"id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
"name": "tool_two_slow",
"result": [("my_key", " slow")],
"error": None,
"interrupts": [],
},
},
{
@@ -5320,6 +5337,8 @@ async def test_branch_then() -> None:
"id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
"name": "finish",
"result": [("my_key", " finished")],
"error": None,
"interrupts": [],
},
},
{
@@ -5363,10 +5382,134 @@ async def test_branch_then() -> None:
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
"my_key": "value prepared",
"market": "DE",
}
assert [
c
async for c in tool_two.astream(
{"my_key": "value", "market": "DE"}, thread1, stream_mode="debug"
)
] == [
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": -1,
"payload": {
"config": {
"tags": [],
"metadata": {"thread_id": "1"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"values": {"my_key": ""},
"metadata": {
"source": "input",
"step": -1,
"writes": {"my_key": "value", "market": "DE"},
},
"next": ["__start__"],
"tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 0,
"payload": {
"config": {
"tags": [],
"metadata": {"thread_id": "1"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"values": {
"my_key": "value",
"market": "DE",
},
"metadata": {
"source": "loop",
"step": 0,
"writes": None,
},
"next": ["prepare"],
"tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
},
},
{
"type": "task",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70",
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70",
"name": "prepare",
"result": [("my_key", " prepared")],
"error": None,
"interrupts": [],
},
},
{
"type": "checkpoint",
"timestamp": AnyStr(),
"step": 1,
"payload": {
"config": {
"tags": [],
"metadata": {"thread_id": "1"},
"callbacks": None,
"recursion_limit": 25,
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
},
},
"values": {
"my_key": "value prepared",
"market": "DE",
},
"metadata": {
"source": "loop",
"step": 1,
"writes": {"prepare": {"my_key": " prepared"}},
},
"next": ["tool_two_slow"],
"tasks": [
{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}
],
},
},
{
"type": "task_result",
"timestamp": AnyStr(),
"step": 2,
"payload": {
"id": "054b0ced-4546-58f4-bee5-548f029e1a8e",
"name": "tool_two_slow",
"result": [],
"error": None,
"interrupts": [{"when": "before", "value": None}],
},
},
]
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value prepared", "market": "DE"},
tasks=(