diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 02d5880a0..433bc231d 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -150,6 +150,7 @@ class PostgresSaver(BasePostgresSaver): } if value["parent_checkpoint_id"] else None, + self._load_writes(value["pending_writes"]), ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: @@ -317,16 +318,6 @@ class PostgresSaver(BasePostgresSaver): task_id (str): Identifier for the task creating the writes. """ with self._cursor(pipeline=True) as cur: - cur.execute( - self.DELETE_WRITES_SQL, - ( - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - len(writes), - ), - ) cur.executemany( self.UPSERT_CHECKPOINT_WRITES_SQL, self._dump_writes( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 7ddb81237..79ae1ddf7 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -135,6 +135,7 @@ class AsyncPostgresSaver(BasePostgresSaver): } if value["parent_checkpoint_id"] else None, + await asyncio.to_thread(self._load_writes, value["pending_writes"]), ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: @@ -273,16 +274,6 @@ class AsyncPostgresSaver(BasePostgresSaver): task_id (str): Identifier for the task creating the writes. """ async with self._cursor(pipeline=True) as cur: - await cur.execute( - self.DELETE_WRITES_SQL, - ( - config["configurable"]["thread_id"], - config["configurable"]["checkpoint_ns"], - config["configurable"]["checkpoint_id"], - task_id, - len(writes), - ), - ) await cur.executemany( self.UPSERT_CHECKPOINT_WRITES_SQL, await asyncio.to_thread( diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 506a13a27..f53bee256 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig from psycopg.types.json import Jsonb from langgraph.checkpoint.base import ( + WRITES_IDX_MAP, BaseCheckpointSaver, Checkpoint, EmptyChannelError, @@ -105,15 +106,6 @@ UPSERT_CHECKPOINT_WRITES_SQL = """ ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING """ -DELETE_WRITES_SQL = """ - DELETE FROM checkpoint_writes - WHERE thread_id = %s - AND checkpoint_ns = %s - AND checkpoint_id = %s - AND task_id = %s - AND idx >= %s -""" - class BasePostgresSaver(BaseCheckpointSaver): SELECT_SQL = SELECT_SQL @@ -121,7 +113,6 @@ class BasePostgresSaver(BaseCheckpointSaver): UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL - DELETE_WRITES_SQL = DELETE_WRITES_SQL jsonplus_serde = JsonPlusSerializer() @@ -210,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver): checkpoint_ns, checkpoint_id, task_id, - idx, + WRITES_IDX_MAP.get(channel, idx), channel, *self.serde.dumps_typed(value), ) diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 9bf4139ad..4767bfb70 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -1,12 +1,13 @@ import sqlite3 import threading -from contextlib import contextmanager +from contextlib import closing, contextmanager from hashlib import md5 from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, @@ -318,7 +319,7 @@ class SqliteSaver(BaseCheckpointSaver): ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" - with self.cursor(transaction=False) as cur: + with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur: cur.execute(query, param_values) for ( thread_id, @@ -329,6 +330,10 @@ class SqliteSaver(BaseCheckpointSaver): checkpoint, metadata, ) in cur: + wcur.execute( + "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + (thread_id, checkpoint_ns, checkpoint_id), + ) yield CheckpointTuple( { "configurable": { @@ -350,6 +355,10 @@ class SqliteSaver(BaseCheckpointSaver): if parent_checkpoint_id else None ), + [ + (task_id, channel, self.serde.loads_typed((type, value))) + for task_id, channel, type, value in wcur + ], ) def put( @@ -424,25 +433,15 @@ class SqliteSaver(BaseCheckpointSaver): task_id (str): Identifier for the task creating the writes. """ with self.lock, self.cursor() as cur: - cur.execute( - "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?", - ( - str(config["configurable"]["thread_id"]), - str(config["configurable"]["checkpoint_ns"]), - str(config["configurable"]["checkpoint_id"]), - task_id, - len(writes), - ), - ) cur.executemany( - "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ ( str(config["configurable"]["thread_id"]), str(config["configurable"]["checkpoint_ns"]), str(config["configurable"]["checkpoint_id"]), task_id, - idx, + WRITES_IDX_MAP.get(channel, idx), channel, *self.serde.dumps_typed(value), ) diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 56d14f613..85869d3f0 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -16,6 +16,7 @@ import aiosqlite from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, @@ -329,14 +330,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver): AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples. """ await self.setup() - where, param_values = search_where(config, filter, before) + where, params = search_where(config, filter, before) query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints {where} ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" - async with self.conn.execute(query, param_values) as cursor: + async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur: async for ( thread_id, checkpoint_ns, @@ -345,7 +346,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver): type, checkpoint, metadata, - ) in cursor: + ) in cur: + await wcur.execute( + "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + (thread_id, checkpoint_ns, checkpoint_id), + ) yield CheckpointTuple( { "configurable": { @@ -367,6 +372,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver): if parent_checkpoint_id else None ), + [ + (task_id, channel, self.serde.loads_typed((type, value))) + async for task_id, channel, type, value in wcur + ], ) async def aput( @@ -433,25 +442,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver): """ await self.setup() async with self.conn.cursor() as cur: - await cur.execute( - "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?", - ( - str(config["configurable"]["thread_id"]), - str(config["configurable"]["checkpoint_ns"]), - str(config["configurable"]["checkpoint_id"]), - task_id, - len(writes), - ), - ) await cur.executemany( - "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ ( str(config["configurable"]["thread_id"]), str(config["configurable"]["checkpoint_ns"]), str(config["configurable"]["checkpoint_id"]), task_id, - idx, + WRITES_IDX_MAP.get(channel, idx), channel, *self.serde.dumps_typed(value), ) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 86c8b0eec..21d021ad5 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -22,6 +22,7 @@ from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( + ERROR, ChannelProtocol, SendProtocol, ) @@ -98,6 +99,7 @@ class Checkpoint(TypedDict): Cleared by the next checkpoint.""" current_tasks: Dict[str, TaskInfo] """Map from task ID to task info.""" + # TODO remove this def empty_checkpoint() -> Checkpoint: @@ -140,6 +142,8 @@ def create_checkpoint( else: values: dict[str, Any] = {} for k, v in channels.items(): + if k not in checkpoint["channel_versions"]: + continue try: values[k] = v.checkpoint() except EmptyChannelError: @@ -437,3 +441,13 @@ def get_checkpoint_id(config: RunnableConfig) -> Optional[str]: return config["configurable"].get( "checkpoint_id", config["configurable"].get("thread_ts") ) + + +""" +Mapping from error type to error index. +Regular writes just map to their index in the list of writes being saved. +Special writes (e.g. errors) map to negative indices, to avoid those writes from +saving regular writes. +Each Checkpointer implementation should use this mapping in put_writes. +""" +WRITES_IDX_MAP = {ERROR: -1} diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 3e8d328e6..52102ffef 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -8,6 +8,7 @@ from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( + WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, Checkpoint, @@ -52,6 +53,9 @@ class MemorySaver( # thread ID -> checkpoint NS -> checkpoint ID -> checkpoint mapping storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]] + writes: defaultdict[ + tuple[str, str, str], dict[tuple[str, int], tuple[str, str, bytes]] + ] def __init__( self, @@ -60,7 +64,7 @@ class MemorySaver( ) -> None: super().__init__(serde=serde) self.storage = defaultdict(lambda: defaultdict(dict)) - self.writes = defaultdict(list) + self.writes = defaultdict(dict) def __enter__(self) -> "MemorySaver": return self @@ -103,7 +107,7 @@ class MemorySaver( if checkpoint_id := get_checkpoint_id(config): if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id): checkpoint, metadata, parent_checkpoint_id = saved - writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)] + writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() return CheckpointTuple( config=config, checkpoint=self.serde.loads_typed(checkpoint), @@ -125,7 +129,7 @@ class MemorySaver( if checkpoints := self.storage[thread_id][checkpoint_ns]: checkpoint_id = max(checkpoints.keys()) checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id] - writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)] + writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() return CheckpointTuple( config={ "configurable": { @@ -204,6 +208,8 @@ class MemorySaver( elif limit is not None: limit -= 1 + writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() + yield CheckpointTuple( config={ "configurable": { @@ -223,6 +229,9 @@ class MemorySaver( } if parent_checkpoint_id else None, + pending_writes=[ + (id, c, self.serde.loads_typed(v)) for id, c, v in writes + ], ) def put( @@ -287,11 +296,10 @@ class MemorySaver( thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"]["checkpoint_ns"] checkpoint_id = config["configurable"]["checkpoint_id"] - key = (thread_id, checkpoint_ns, checkpoint_id) - self.writes[key] = [w for w in self.writes[key] if w[0] != task_id] - self.writes[key].extend( - [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes] - ) + outer_key = (thread_id, checkpoint_ns, checkpoint_id) + for idx, (c, v) in enumerate(writes): + inner_key = (task_id, WRITES_IDX_MAP.get(c, idx)) + self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v)) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Asynchronous version of get_tuple. diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index de61ef78b..d166a6a0d 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -12,6 +12,8 @@ from typing import ( from langchain_core.runnables import RunnableConfig from typing_extensions import Self +ERROR = "__error__" + Value = TypeVar("Value") Update = TypeVar("Update") C = TypeVar("C") diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 00dc36e58..27a7689fa 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Sequence from langgraph.checkpoint.base import EmptyChannelError from langgraph.constants import Interrupt @@ -32,7 +32,7 @@ class InvalidUpdateError(Exception): class GraphInterrupt(Exception): """Raised when a subgraph is interrupted.""" - def __init__(self, interrupts: list[Interrupt]) -> None: + def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None: super().__init__(interrupts) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 270a3c7ba..c06f2bce3 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -65,16 +65,10 @@ from langgraph.constants import ( CONFIG_KEY_SEND, ERROR, INTERRUPT, - Interrupt, ) from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ManagedValueSpec -from langgraph.pregel.algo import ( - apply_writes, - local_read, - prepare_next_tasks, - should_interrupt, -) +from langgraph.pregel.algo import apply_writes, local_read, prepare_next_tasks from langgraph.pregel.debug import ( print_step_checkpoint, print_step_tasks, @@ -570,7 +564,7 @@ class Pregel( # update channels with ChannelsManager(self.channels, checkpoint, config) as ( channels, - managed, + _, ): # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() @@ -606,31 +600,6 @@ class Pregel( checkpoint, channels, [task], self.checkpointer.get_next_version ), "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) - # check interrupt before - if tasks := should_interrupt( - checkpoint, - self.interrupt_before_nodes, - prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - step + 2, - for_execution=False, - ), - ): - for t in tasks: - self.checkpointer.put_writes( - { - "configurable": { - **checkpoint_config["configurable"], - "checkpoint_id": checkpoint["id"], - } - }, - [(INTERRUPT, Interrupt("before"))], - t.id, - ) return self.checkpointer.put( checkpoint_config, checkpoint, @@ -713,7 +682,7 @@ class Pregel( # update channels, acting as the chosen node async with AsyncChannelsManager(self.channels, checkpoint, config) as ( channels, - managed, + _, ): # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() @@ -749,35 +718,6 @@ class Pregel( checkpoint, channels, [task], self.checkpointer.get_next_version ), "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) - # check interrupt before - if tasks := should_interrupt( - checkpoint, - self.interrupt_before_nodes, - prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - step + 2, - for_execution=False, - ), - ): - await asyncio.gather( - *( - self.checkpointer.aput_writes( - { - "configurable": { - **checkpoint_config["configurable"], - "checkpoint_id": checkpoint["id"], - } - }, - [(INTERRUPT, Interrupt("before"))], - t.id, - ) - for t in tasks - ) - ) return await self.checkpointer.aput( checkpoint_config, checkpoint, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 061b932c0..a501232a9 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -41,7 +41,6 @@ from langgraph.constants import ( ERROR, INPUT, INTERRUPT, - Interrupt, ) from langgraph.errors import EmptyInputError, GraphInterrupt from langgraph.managed.base import ( @@ -162,6 +161,8 @@ class PregelLoop: def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: """Put writes for a task, to be read by the next tick.""" + if not writes: + return self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) if self.checkpointer_put_writes is not None: self.submit( @@ -238,13 +239,10 @@ class PregelLoop: } ) # after execution, check if we should interrupt - if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks): + if should_interrupt(self.checkpoint, interrupt_after, self.tasks): self.status = "interrupt_after" - interrupts = [(t.id, Interrupt("after")) for t in tasks] - for tid, interrupt in interrupts: - self.put_writes(tid, [(INTERRUPT, interrupt)]) if self.is_nested: - raise GraphInterrupt([i[1] for i in interrupts]) + raise GraphInterrupt() else: return False else: @@ -308,13 +306,10 @@ class PregelLoop: ) # before execution, check if we should interrupt - if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks): + if should_interrupt(self.checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" - interrupts = [(t.id, Interrupt("before")) for t in tasks] - for tid, interrupt in interrupts: - self.put_writes(tid, [(INTERRUPT, interrupt)]) if self.is_nested: - raise GraphInterrupt([i[1] for i in interrupts]) + raise GraphInterrupt() else: return False diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index 9aa030d3b..a98962cdc 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -1,3 +1,6 @@ +from typing import Any, Sequence + + class AnyStr(str): def __init__(self) -> None: super().__init__() @@ -9,6 +12,17 @@ class AnyStr(str): return hash(str(self)) +class AnyVersion: + def __init__(self) -> None: + super().__init__() + + def __eq__(self, other: object) -> bool: + return isinstance(other, (str, int, float)) + + def __hash__(self) -> int: + return hash(str(self)) + + class ExceptionLike: def __init__(self, exc: Exception) -> None: self.exc = exc @@ -22,3 +36,24 @@ class ExceptionLike: def __hash__(self) -> int: return hash((self.exc.__class__, str(self.exc))) + + def __repr__(self) -> str: + return str(self.exc) + + +class UnsortedSequence: + def __init__(self, *values: Any) -> None: + self.seq = values + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, Sequence) + and len(self.seq) == len(value) + and all(a in value for a in self.seq) + ) + + def __hash__(self) -> int: + return hash(frozenset(self.seq)) + + def __repr__(self) -> str: + return repr(self.seq) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 3e40f31fc..56e7ed5a3 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -67,7 +67,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore -from tests.any_str import AnyStr, ExceptionLike +from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -1411,6 +1411,147 @@ def test_pending_writes_resume( # both the pending write and the new write were applied, 1 + 2 + 3 = 6 assert graph.invoke(None, thread1) == {"value": 6} + # check all final checkpoints + checkpoints = [c for c in checkpointer.list(thread1)] + # we should have 3 + assert len(checkpoints) == 3 + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "current_tasks": {}, + "pending_sends": [], + "versions_seen": { + "one": { + "start:one": AnyVersion(), + }, + "two": { + "start:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + }, + "channel_versions": { + "one": AnyVersion(), + "two": AnyVersion(), + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + "channel_values": {"one": "one", "two": "two", "value": 6}, + }, + metadata={ + "step": 1, + "source": "loop", + "writes": {"one": {"value": 2}, "two": {"value": 3}}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "current_tasks": {}, + "pending_sends": [], + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "start:one": "__start__", + "start:two": "__start__", + }, + }, + metadata={"step": 0, "source": "loop", "writes": None}, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=UnsortedSequence( + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + (AnyStr(), "__error__", ExceptionLike(ConnectionError("I'm not good"))), + (AnyStr(), "two", "two"), + (AnyStr(), "value", 3), + ), + ) + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "current_tasks": {}, + "pending_sends": [], + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), + }, + "channel_values": {"__start__": {"value": 1}}, + }, + metadata={"step": -1, "source": "input", "writes": {"value": 1}}, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "start:one", "__start__"), + (AnyStr(), "start:two", "__start__"), + ), + ) + def test_cond_edge_after_send() -> None: class Node: @@ -2235,7 +2376,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2281,7 +2422,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "input": "what is weather in sf", }, }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2442,7 +2583,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3043,7 +3184,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3083,7 +3224,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3192,7 +3333,7 @@ def test_conditional_state_graph( values={ "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3217,7 +3358,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3270,7 +3411,7 @@ def test_conditional_state_graph( ) ], }, - tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4872,13 +5013,7 @@ def test_message_graph( id="ai1", ), ], - tasks=( - PregelTask( - AnyStr(), - "tools", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4923,7 +5058,7 @@ def test_message_graph( ], ), ], - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5005,7 +5140,7 @@ def test_message_graph( id="ai2", ), ], - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5597,13 +5732,7 @@ def test_root_graph( id="ai1", ), ], - tasks=( - PregelTask( - AnyStr(), - "tools", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5648,7 +5777,7 @@ def test_root_graph( ], ), ], - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5730,7 +5859,7 @@ def test_root_graph( id="ai2", ), ], - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -6340,13 +6469,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_slow", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6380,13 +6503,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6420,13 +6537,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6437,13 +6548,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "valuekey", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6744,13 +6849,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_slow", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6788,13 +6887,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6841,7 +6934,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slow", "market": "DE", }, - tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "finish"),), next=("finish",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6860,7 +6953,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slower", "market": "DE", }, - tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "finish"),), next=("finish",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -7117,7 +7210,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], }, - tasks=(PregelTask(AnyStr(), "qa", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "qa"),), next=("qa",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 692d101e0..0308deb54 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -62,7 +62,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore -from tests.any_str import AnyStr, ExceptionLike +from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -207,7 +207,15 @@ async def test_node_cancellation_on_other_node_exception() -> None: assert inner_task_cancelled -async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_dynamic_interrupt( + checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -245,51 +253,49 @@ async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: "market": "US", } - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile(checkpointer=saver) + tool_two = tool_two_graph.compile(checkpointer=checkpointer) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - 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 ⛰️", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value ⛰️", "market": "DE"}, - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - interrupts=(Interrupt("during", "Just because..."),), - ), + 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 ⛰️", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("during", "Just because..."),), ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) + # TODO use aget_state_history @pytest.mark.parametrize( @@ -1609,6 +1615,10 @@ async def test_pending_writes_resume( error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) assert error_write[0] != non_error_writes[0][0] + # TODO arguably this shouldn't even run the failed task again, + # and should require empty update_state (ie new checkpoint_id) + # in order to try again + # resume execution with pytest.raises(ValueError, match="I'm not good"): await graph.ainvoke(None, thread1) @@ -1627,6 +1637,147 @@ async def test_pending_writes_resume( # both the pending write and the new write were applied, 1 + 2 + 3 = 6 assert await graph.ainvoke(None, thread1) == {"value": 6} + # check all final checkpoints + checkpoints = [c async for c in checkpointer.alist(thread1)] + # we should have 3 + assert len(checkpoints) == 3 + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "current_tasks": {}, + "pending_sends": [], + "versions_seen": { + "one": { + "start:one": AnyVersion(), + }, + "two": { + "start:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + }, + "channel_versions": { + "one": AnyVersion(), + "two": AnyVersion(), + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + "channel_values": {"one": "one", "two": "two", "value": 6}, + }, + metadata={ + "step": 1, + "source": "loop", + "writes": {"one": {"value": 2}, "two": {"value": 3}}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "current_tasks": {}, + "pending_sends": [], + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "start:one": "__start__", + "start:two": "__start__", + }, + }, + metadata={"step": 0, "source": "loop", "writes": None}, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=UnsortedSequence( + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + (AnyStr(), "__error__", ExceptionLike(ValueError("I'm not good"))), + (AnyStr(), "two", "two"), + (AnyStr(), "value", 3), + ), + ) + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "current_tasks": {}, + "pending_sends": [], + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), + }, + "channel_values": {"__start__": {"value": 1}}, + }, + metadata={"step": -1, "source": "input", "writes": {"value": 1}}, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "start:one", "__start__"), + (AnyStr(), "start:two", "__start__"), + ), + ) + async def test_cond_edge_after_send() -> None: class Node: @@ -2516,13 +2667,7 @@ async def test_conditional_graph() -> None: ), }, }, - tasks=( - PregelTask( - AnyStr(), - "tools", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2572,13 +2717,7 @@ async def test_conditional_graph() -> None: "input": "what is weather in sf", }, }, - tasks=( - PregelTask( - AnyStr(), - "tools", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2750,13 +2889,7 @@ async def test_conditional_graph() -> None: ), }, }, - tasks=( - PregelTask( - AnyStr(), - "tools", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3324,13 +3457,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, - tasks=( - PregelTask( - AnyStr(), - "tools", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3374,7 +3501,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, - tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),), + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4957,13 +5084,7 @@ async def test_start_branch_then() -> None: ] assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value", "market": "DE"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_slow", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5005,13 +5126,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -5053,13 +5168,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "value", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -5074,13 +5183,7 @@ async def test_start_branch_then() -> None: await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "valuekey", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -5497,28 +5600,10 @@ async def test_branch_then() -> None: ], }, }, - { - "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=( - PregelTask( - AnyStr(), - "tool_two_slow", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5564,13 +5649,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, - tasks=( - PregelTask( - AnyStr(), - "tool_two_fast", - interrupts=(Interrupt("before"),), - ), - ), + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[