WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 17:57:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix semantics of put_writes/list (#1436)

Browse Source 
* Fix semantics of put_writes/list

- put_writes(error) should not prevent saving future successful if task is retried successfully
- put_writes(writes) should be a no-op if non-error writes already exist for that task (this prevents tasks executed more than once from modifying writes previously saved / acted on)
- checkpoints should not include channel default values (ie those without a version)
- list() should fetch and return writes for each checkpoint

* Lint

* Rm print

* Fix import

* Lint
This commit is contained in:
Nuno Campos
2024-08-22 19:04:42 +00:00
committed by GitHub
parent 38daba5259
commit a261e1a497
14 changed files with 480 additions and 343 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+1 -10
View File
@@ -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(
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+1 -10
View File
@@ -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(
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+2 -11
View File
@@ -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),
)
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+13 -14
View File
@@ -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),
)
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+14 -15
View File
@@ -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),
)
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+14
View File
@@ -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}
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+16 -8
View File
@@ -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.
libs/checkpoint/langgraph/checkpoint/serde/types.py
+2
View File
@@ -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")
libs/langgraph/langgraph/errors.py
+2 -2
View File
@@ -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)
libs/langgraph/langgraph/pregel/__init__.py
+3 -63
View File
@@ -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,
libs/langgraph/langgraph/pregel/loop.py
+6 -11
View File
@@ -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
libs/langgraph/tests/any_str.py
+35
View File
@@ -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)
libs/langgraph/tests/test_pregel.py
+165 -72
View File
@@ -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"],
libs/langgraph/tests/test_pregel_async.py
+206 -127
View File
@@ -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[