From d85e267a83bec2e080eae26e1e3486c39f8020d1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 21 Aug 2024 16:23:48 -0700 Subject: [PATCH] Add interrupts and errors to debug stream mode - move remaining output code from Pregel.stream to PregelLoop --- libs/langgraph/langgraph/pregel/__init__.py | 122 +++++++--------- libs/langgraph/langgraph/pregel/debug.py | 16 +- libs/langgraph/langgraph/pregel/io.py | 19 ++- libs/langgraph/langgraph/pregel/loop.py | 36 ++++- libs/langgraph/tests/test_pregel.py | 23 ++- libs/langgraph/tests/test_pregel_async.py | 153 +++++++++++++++++++- 6 files changed, 272 insertions(+), 97 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6e4d3e002..563f097f8 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -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: diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 5102cd8c5..72b6cdbe5 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -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) ], diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 53e77557e..6c9205248 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -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(): diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index aa256b2f7..11b937eea 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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() diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 74897efcc..c1fc68c2a 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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": [], }, }, { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index ef78c8bdb..692d101e0 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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=(