From 0d13c6b159a609ecaf2ec69a5b63e1b69879c2c6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Mar 2024 17:22:39 -0700 Subject: [PATCH] Add history tracking to in-memory, sqlite and aiosqlite checkpointers - Add Pregel.get_state_history and .aget_state_history methods to get history iterator - Update checkpointer base class with new list and get_tuple methods - Rewrite in-memory checkpointer class to track history - Rewrite sqlite and aiosqlite checkpointers to track history - Add new tests for history tracking --- langgraph/checkpoint/aiosqlite.py | 122 ++++++++++++++++++------------ langgraph/checkpoint/base.py | 66 +++++++++++++--- langgraph/checkpoint/memory.py | 55 +++++++++----- langgraph/checkpoint/sqlite.py | 99 +++++++++++++++++------- langgraph/pregel/__init__.py | 70 +++++++++++++++-- tests/memory_assert.py | 14 +++- tests/test_pregel.py | 118 +++++++++++++++++++++-------- tests/test_pregel_async.py | 122 +++++++++++++++++++++--------- 8 files changed, 486 insertions(+), 180 deletions(-) diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index 580a3a102..199a23c6c 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -1,15 +1,16 @@ import pickle -from typing import Optional +from contextlib import AbstractAsyncContextManager +from types import TracebackType +from typing import AsyncIterator, Optional, Self import aiosqlite from langchain_core.pydantic_v1 import Field from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.utils import ConfigurableFieldSpec -from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint +from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple -class AsyncSqliteSaver(BaseCheckpointSaver): +class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): conn: aiosqlite.Connection is_setup: bool = Field(False, init=False, repr=False) @@ -21,65 +22,92 @@ class AsyncSqliteSaver(BaseCheckpointSaver): def from_conn_string(cls, conn_string: str) -> "AsyncSqliteSaver": return AsyncSqliteSaver(conn=aiosqlite.connect(conn_string)) - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id="thread_id", - annotation=str, - name="Thread ID", - description=None, - default="", - is_shared=True, - ), - ] + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + __exc_type: type[BaseException] | None, + __exc_value: BaseException | None, + __traceback: TracebackType | None, + ) -> bool | None: + return await self.conn.close() async def setup(self) -> None: - print("hello") if self.is_setup: return - try: - await self.conn - await self.conn.executescript( - """ - CREATE TABLE IF NOT EXISTS checkpoints ( - thread_id TEXT PRIMARY KEY, - checkpoint BLOB - ); - """ - ) + await self.conn + async with self.conn.executescript( + """ + CREATE TABLE IF NOT EXISTS checkpoints ( + thread_id TEXT NOT NULL, + thread_ts TEXT NOT NULL, + checkpoint BLOB, + PRIMARY KEY (thread_id, thread_ts) + ); + """ + ): await self.conn.commit() - print("good bye") + self.is_setup = True - self.is_setup = True - except BaseException as e: - print(e) - raise e + async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + await self.setup() + if config["configurable"].get("thread_ts"): + async with self.conn.execute( + "SELECT checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", + ( + config["configurable"]["thread_id"], + config["configurable"]["thread_ts"], + ), + ) as cursor: + if value := await cursor.fetchone(): + return CheckpointTuple(config, pickle.loads(value[0])) + else: + async with self.conn.execute( + "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1", + (config["configurable"]["thread_id"],), + ) as cursor: + if value := await cursor.fetchone(): + return CheckpointTuple( + { + "configurable": { + "thread_id": value[0], + "thread_ts": value[1], + } + }, + pickle.loads(value[2]), + ) - def get(self, config: RunnableConfig) -> Optional[Checkpoint]: - raise NotImplementedError - - def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: - raise NotImplementedError - - async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: + async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]: await self.setup() async with self.conn.execute( - "SELECT checkpoint FROM checkpoints WHERE thread_id = ?", + "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC", (config["configurable"]["thread_id"],), ) as cursor: - if value := await cursor.fetchone(): - return pickle.loads(value[0]) + async for thread_id, thread_ts, value in cursor: + yield CheckpointTuple( + {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, + pickle.loads(value), + ) - async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: + async def aput( + self, config: RunnableConfig, checkpoint: Checkpoint + ) -> RunnableConfig: await self.setup() - await self.conn.execute( - "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)", + async with self.conn.execute( + "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, checkpoint) VALUES (?, ?, ?)", ( config["configurable"]["thread_id"], + checkpoint["ts"], pickle.dumps(checkpoint), ), - ) - await self.conn.commit() + ): + await self.conn.commit() + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "thread_ts": checkpoint["ts"], + } + } diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 99b0e294a..b5000a303 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -1,9 +1,9 @@ import asyncio -from abc import ABC, abstractmethod +from abc import ABC from collections import defaultdict from copy import deepcopy from datetime import datetime, timezone -from typing import Any, Optional, TypedDict +from typing import Any, AsyncIterator, Iterator, NamedTuple, Optional, TypedDict from langchain_core.load.serializable import Serializable from langchain_core.runnables import RunnableConfig @@ -49,25 +49,71 @@ class CheckpointAt(StrEnum): END_OF_RUN = "end_of_run" +class CheckpointTuple(NamedTuple): + config: RunnableConfig + checkpoint: Checkpoint + + +CheckpointThreadId = ConfigurableFieldSpec( + id="thread_id", + annotation=str, + name="Thread ID", + description=None, + default="", + is_shared=True, +) + +CheckpointThreadTs = ConfigurableFieldSpec( + id="thread_ts", + annotation=Optional[str], + name="Thread Timestamp", + description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", + default=None, + is_shared=True, +) + + class BaseCheckpointSaver(Serializable, ABC): at: CheckpointAt = CheckpointAt.END_OF_RUN @property def config_specs(self) -> list[ConfigurableFieldSpec]: - return [] + return [CheckpointThreadId, CheckpointThreadTs] - @abstractmethod def get(self, config: RunnableConfig) -> Optional[Checkpoint]: - ... + if value := self.get_tuple(config): + return value.checkpoint - @abstractmethod - def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: - ... + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + raise NotImplementedError + + def list(self, config: RunnableConfig) -> Iterator[CheckpointTuple]: + raise NotImplementedError + + def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig: + raise NotImplementedError async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: - return await asyncio.get_running_loop().run_in_executor(None, self.get, config) + if value := await self.aget_tuple(config): + return value.checkpoint - async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: + async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + return await asyncio.get_running_loop().run_in_executor( + None, self.get_tuple, config + ) + + async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]: + loop = asyncio.get_running_loop() + iter = loop.run_in_executor(None, self.list, config) + while True: + try: + yield await loop.run_in_executor(None, next, iter) + except StopIteration: + return + + async def aput( + self, config: RunnableConfig, checkpoint: Checkpoint + ) -> RunnableConfig: return await asyncio.get_running_loop().run_in_executor( None, self.put, config, checkpoint ) diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 7812c6f9a..059c58573 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -1,30 +1,47 @@ +from collections import defaultdict from typing import Optional from langchain_core.pydantic_v1 import Field from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.utils import ConfigurableFieldSpec -from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint +from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple class MemorySaver(BaseCheckpointSaver): - storage: dict[str, Checkpoint] = Field(default_factory=dict) - - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id="thread_id", - annotation=str, - name="Thread ID", - description=None, - default="", - is_shared=True, - ), - ] + storage: defaultdict[str, dict[str, Checkpoint]] = Field( + default_factory=lambda: defaultdict(dict) + ) def get(self, config: RunnableConfig) -> Optional[Checkpoint]: - return self.storage.get(config["configurable"]["thread_id"], None) + if value := self.get_tuple(config): + return value.checkpoint - def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: - return self.storage.update({config["configurable"]["thread_id"]: checkpoint}) + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + if config["configurable"].get("thread_ts"): + if checkpoint := self.storage[config["configurable"]["thread_id"]].get( + config["configurable"]["thread_ts"] + ): + return CheckpointTuple(config=config, checkpoint=checkpoint) + else: + if checkpoints := self.storage[config["configurable"]["thread_id"]]: + thread_ts = max(checkpoints.keys()) + return CheckpointTuple( + config={ + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "thread_ts": thread_ts, + } + }, + checkpoint=checkpoints[thread_ts], + ) + + def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig: + self.storage[config["configurable"]["thread_id"]].update( + {checkpoint["ts"]: checkpoint} + ) + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "thread_ts": checkpoint["ts"], + } + } diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index 75958708e..f73ebd840 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -1,16 +1,16 @@ import pickle import sqlite3 -from contextlib import contextmanager -from typing import Optional +from contextlib import AbstractContextManager, contextmanager +from types import TracebackType +from typing import AsyncIterator, Iterator, Optional, Self from langchain_core.pydantic_v1 import Field from langchain_core.runnables import RunnableConfig -from langchain_core.runnables.utils import ConfigurableFieldSpec -from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint +from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple -class SqliteSaver(BaseCheckpointSaver): +class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): conn: sqlite3.Connection is_setup: bool = Field(False, init=False, repr=False) @@ -22,18 +22,16 @@ class SqliteSaver(BaseCheckpointSaver): def from_conn_string(cls, conn_string: str) -> "SqliteSaver": return SqliteSaver(conn=sqlite3.connect(conn_string)) - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id="thread_id", - annotation=str, - name="Thread ID", - description=None, - default="", - is_shared=True, - ), - ] + def __enter__(self) -> Self: + return self + + def __exit__( + self, + __exc_type: type[BaseException] | None, + __exc_value: BaseException | None, + __traceback: TracebackType | None, + ) -> bool | None: + return self.conn.close() def setup(self) -> None: if self.is_setup: @@ -42,8 +40,10 @@ class SqliteSaver(BaseCheckpointSaver): self.conn.executescript( """ CREATE TABLE IF NOT EXISTS checkpoints ( - thread_id TEXT PRIMARY KEY, - checkpoint BLOB + thread_id TEXT NOT NULL, + thread_ts TEXT NOT NULL, + checkpoint BLOB, + PRIMARY KEY (thread_id, thread_ts) ); """ ) @@ -61,27 +61,72 @@ class SqliteSaver(BaseCheckpointSaver): self.conn.commit() cur.close() - def get(self, config: RunnableConfig) -> Optional[Checkpoint]: + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + with self.cursor(transaction=False) as cur: + if config["configurable"].get("thread_ts"): + cur.execute( + "SELECT checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", + ( + config["configurable"]["thread_id"], + config["configurable"]["thread_ts"], + ), + ) + if value := cur.fetchone(): + return CheckpointTuple(config, pickle.loads(value[0])) + else: + cur.execute( + "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1", + (config["configurable"]["thread_id"],), + ) + if value := cur.fetchone(): + return CheckpointTuple( + { + "configurable": { + "thread_id": value[0], + "thread_ts": value[1], + } + }, + pickle.loads(value[2]), + ) + + def list(self, config: RunnableConfig) -> Iterator[CheckpointTuple]: with self.cursor(transaction=False) as cur: cur.execute( - "SELECT checkpoint FROM checkpoints WHERE thread_id = ?", + "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC", (config["configurable"]["thread_id"],), ) - if value := cur.fetchone(): - return pickle.loads(value[0]) + for thread_id, thread_ts, value in cur: + yield CheckpointTuple( + {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, + pickle.loads(value), + ) - def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: + def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig: with self.cursor() as cur: cur.execute( - "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)", + "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, checkpoint) VALUES (?, ?, ?)", ( config["configurable"]["thread_id"], + checkpoint["ts"], pickle.dumps(checkpoint), ), ) + return { + "configurable": { + "thread_id": config["configurable"]["thread_id"], + "thread_ts": checkpoint["ts"], + } + } - async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]: + async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: raise NotImplementedError - async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: + async def alist( + self, config: RunnableConfig + ) -> AsyncIterator[tuple[RunnableConfig, Checkpoint]]: + raise NotImplementedError + + async def aput( + self, config: RunnableConfig, checkpoint: Checkpoint + ) -> RunnableConfig: raise NotImplementedError diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5ad9c1922..5fe4fd2aa 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -163,6 +163,8 @@ class StateSnapshot(NamedTuple): """Current values of channels""" next: tuple[str] """Nodes to execute in the next step, if any""" + config: RunnableConfig + """Config used to fetch this snapshot""" class Pregel( @@ -278,8 +280,9 @@ class Pregel( if not self.checkpointer: raise ValueError("No checkpointer set") - checkpoint = self.checkpointer.get(config) - checkpoint = checkpoint or empty_checkpoint() + saved = self.checkpointer.get_tuple(config) + checkpoint = saved.checkpoint if saved else empty_checkpoint() + config = saved.config if saved else config with ChannelsManager(self.channels, checkpoint) as channels: _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, update_seen=False @@ -294,14 +297,16 @@ class Pregel( if isinstance(self.snapshot_channels, str) else values, tuple(name for _, _, name in next_tasks), + config, ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: if not self.checkpointer: raise ValueError("No checkpointer set") - checkpoint = await self.checkpointer.aget(config) - checkpoint = checkpoint or empty_checkpoint() + saved = await self.checkpointer.aget_tuple(config) + checkpoint = saved.checkpoint if saved else empty_checkpoint() + config = saved.config if saved else config async with AsyncChannelsManager(self.channels, checkpoint) as channels: _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, update_seen=False @@ -316,11 +321,58 @@ class Pregel( if isinstance(self.snapshot_channels, str) else values, tuple(name for _, _, name in next_tasks), + config, ) + def get_state_history(self, config: RunnableConfig) -> Iterator[StateSnapshot]: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + for config, checkpoint in self.checkpointer.list(config): + with ChannelsManager(self.channels, checkpoint) as channels: + _, next_tasks = _prepare_next_tasks( + checkpoint, self.nodes, channels, update_seen=False + ) + values = { + k: _read_channel(channels, k) + for k in channels + if k in self.snapshot_channels_list + } + yield StateSnapshot( + values[self.snapshot_channels] + if isinstance(self.snapshot_channels, str) + else values, + tuple(name for _, _, name in next_tasks), + config, + ) + + async def aget_state_history( + self, config: RunnableConfig + ) -> AsyncIterator[StateSnapshot]: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + async for config, checkpoint in self.checkpointer.alist(config): + async with AsyncChannelsManager(self.channels, checkpoint) as channels: + _, next_tasks = _prepare_next_tasks( + checkpoint, self.nodes, channels, update_seen=False + ) + values = { + k: _read_channel(channels, k) + for k in channels + if k in self.snapshot_channels_list + } + yield StateSnapshot( + values[self.snapshot_channels] + if isinstance(self.snapshot_channels, str) + else values, + tuple(name for _, _, name in next_tasks), + config, + ) + def update_state( self, config: RunnableConfig, values: dict[str, Any] | Any - ) -> None: + ) -> RunnableConfig: if not self.checkpointer: raise ValueError("No checkpointer set") @@ -338,11 +390,13 @@ class Pregel( for k in self.snapshot_channels_list: version = checkpoint["channel_versions"][k] checkpoint["versions_seen"][INTERRUPT][k] = version - self.checkpointer.put(config, create_checkpoint(checkpoint, channels)) + return self.checkpointer.put( + config, create_checkpoint(checkpoint, channels) + ) async def aupdate_state( self, config: RunnableConfig, values: dict[str, Any] | Any - ) -> None: + ) -> RunnableConfig: if not self.checkpointer: raise ValueError("No checkpointer set") @@ -360,7 +414,7 @@ class Pregel( for k in self.snapshot_channels or self.channels: version = checkpoint["channel_versions"][k] checkpoint["versions_seen"][INTERRUPT][k] = version - await self.checkpointer.aput( + return await self.checkpointer.aput( config, create_checkpoint(checkpoint, channels) ) diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 4362cf939..a2d1c9ee8 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -1,3 +1,5 @@ +from collections import defaultdict + from langchain_core.pydantic_v1 import Field from langgraph.checkpoint.base import Checkpoint, CheckpointAt, copy_checkpoint @@ -5,15 +7,19 @@ from langgraph.checkpoint.memory import MemorySaver class MemorySaverAssertImmutable(MemorySaver): - storage_for_copies: dict[str, Checkpoint] = Field(default_factory=dict) + storage_for_copies: defaultdict[str, dict[str, Checkpoint]] = Field( + default_factory=lambda: defaultdict(dict) + ) at = CheckpointAt.END_OF_STEP - def put(self, config: dict, checkpoint: dict) -> None: + def put(self, config: dict, checkpoint: Checkpoint) -> None: # assert checkpoint hasn't been modified since last written thread_id = config["configurable"]["thread_id"] if saved := super().get(config): - assert self.storage_for_copies[thread_id] == saved - self.storage_for_copies[thread_id] = copy_checkpoint(checkpoint) + assert self.storage_for_copies[thread_id][saved["ts"]] == saved + self.storage_for_copies[thread_id][checkpoint["ts"]] = copy_checkpoint( + checkpoint + ) # call super to write checkpoint super().put(config, checkpoint) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index a85fa080e..9558dbe60 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -503,39 +503,79 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = SqliteSaver.from_conn_string(":memory:") + with SqliteSaver.from_conn_string(":memory:") as memory: + app = Pregel( + nodes={"one": one}, + channels={"total": BinaryOperatorAggregate(int, operator.add)}, + checkpointer=memory, + ) - app = Pregel( - nodes={"one": one}, - channels={"total": BinaryOperatorAggregate(int, operator.add)}, - checkpointer=memory, - ) + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert app.invoke(2, thread_1) == 2 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"] + # total is now 2, so output is 2+3=5 + assert app.invoke(3, thread_1) == 5 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"] + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + app.invoke(4, thread_1) + # checkpoint is not updated + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 - # total starts out as 0, so output is 0+2=2 - assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 2 - # total is now 2, so output is 2+3=5 - assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - app.invoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # on a new thread, total starts out as 0, so output is 0+5=5 - assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - checkpoint = memory.get({"configurable": {"thread_id": "2"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 5 + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert app.invoke(5, thread_2) == 5 + state = app.get_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 7 + state = app.get_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + + # list all checkpoints for thread 1 + thread_1_history = [c for c in app.get_state_history(thread_1)] + # there are 2: one for each successful ainvoke() + assert len(thread_1_history) == 2 + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["thread_ts"] + > thread_1_history[1].config["configurable"]["thread_ts"] + ) + # the second checkpoint + assert thread_1_history[0].values["total"] == 7 + # the first checkpoint + assert thread_1_history[1].values["total"] == 2 + # can get each checkpoint using aget with config + assert ( + memory.get(thread_1_history[0].config)["ts"] + == thread_1_history[0].config["configurable"]["thread_ts"] + ) + assert ( + memory.get(thread_1_history[1].config)["ts"] + == thread_1_history[1].config["configurable"]["thread_ts"] + ) + + thread_1_next_config = app.update_state( + thread_1_history[1].config, {"total": 10} + ) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["thread_ts"] + > thread_1_history[0].config["configurable"]["thread_ts"] + ) + # 1 more checkpoint in history + assert len(list(app.get_state_history(thread_1))) == 3 + # the latest checkpoint is the updated one + assert app.get_state(thread_1) == app.get_state(thread_1_next_config) def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -942,6 +982,13 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": None, }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + ) + assert ( + app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ + "thread_ts" + ] + is not None ) app_w_interrupt.update_state( @@ -971,6 +1018,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": None, }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1088,6 +1136,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": None, }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) app_w_interrupt.update_state( @@ -1117,6 +1166,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": None, }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1234,6 +1284,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": None, }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1571,6 +1622,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) app_w_interrupt.update_state( @@ -1595,6 +1647,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -1686,6 +1739,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) app_w_interrupt.update_state( @@ -1710,6 +1764,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) assert [c for c in app_w_interrupt.stream(None, config)] == [ @@ -2478,6 +2533,7 @@ def test_message_graph(snapshot: SnapshotAssertion) -> None: ), ], next=("agent:edges",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, ) # TODO use update_state once we have message ids diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index eb06b9e74..0fbb65994 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -502,7 +502,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: assert checkpoint["channel_values"].get("total") == 5 -async def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: +async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -517,42 +517,86 @@ async def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = AsyncSqliteSaver.from_conn_string(":memory:") + async with AsyncSqliteSaver.from_conn_string(":memory:") as memory: + app = Pregel( + nodes={"one": one}, + channels={"total": BinaryOperatorAggregate(int, operator.add)}, + checkpointer=memory, + debug=True, + ) - app = Pregel( - nodes={"one": one}, - channels={"total": BinaryOperatorAggregate(int, operator.add)}, - checkpointer=memory, - debug=True, - ) + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, thread_1) == 2 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert ( + state.config["configurable"]["thread_ts"] + == (await memory.aget(thread_1))["ts"] + ) + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, thread_1) == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["thread_ts"] + == (await memory.aget(thread_1))["ts"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, thread_1) + # checkpoint is not updated + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 2 - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - checkpoint = await memory.aget({"configurable": {"thread_id": "2"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 5 + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, thread_2) == 5 + state = await app.aget_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 7 + state = await app.aget_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 - await memory.conn.close() + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 2: one for each successful ainvoke() + assert len(thread_1_history) == 2 + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["thread_ts"] + > thread_1_history[1].config["configurable"]["thread_ts"] + ) + # the second checkpoint + assert thread_1_history[0].values["total"] == 7 + # the first checkpoint + assert thread_1_history[1].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await memory.aget(thread_1_history[0].config))[ + "ts" + ] == thread_1_history[0].config["configurable"]["thread_ts"] + assert (await memory.aget(thread_1_history[1].config))[ + "ts" + ] == thread_1_history[1].config["configurable"]["thread_ts"] + + thread_1_next_config = await app.aupdate_state( + thread_1_history[1].config, {"total": 10} + ) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["thread_ts"] + > thread_1_history[0].config["configurable"]["thread_ts"] + ) + # 1 more checkpoint in history + assert len([h async for h in app.aget_state_history(thread_1)]) == 3 + # the latest checkpoint is the updated one + assert await app.aget_state(thread_1) == await app.aget_state( + thread_1_next_config + ) async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -988,6 +1032,7 @@ async def test_conditional_graph() -> None: "tools": None, }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) await app_w_interrupt.aupdate_state( @@ -1017,6 +1062,7 @@ async def test_conditional_graph() -> None: "tools": None, }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1137,6 +1183,7 @@ async def test_conditional_graph() -> None: "tools": None, }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) await app_w_interrupt.aupdate_state( @@ -1166,6 +1213,7 @@ async def test_conditional_graph() -> None: "tools": None, }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1286,6 +1334,7 @@ async def test_conditional_graph() -> None: "tools": None, }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1621,6 +1670,7 @@ async def test_conditional_graph_state() -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) await app_w_interrupt.aupdate_state( @@ -1645,6 +1695,7 @@ async def test_conditional_graph_state() -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -1737,6 +1788,7 @@ async def test_conditional_graph_state() -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) await app_w_interrupt.aupdate_state( @@ -1761,6 +1813,7 @@ async def test_conditional_graph_state() -> None: "intermediate_steps": [], }, next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ @@ -2512,6 +2565,7 @@ async def test_message_graph() -> None: ), ], next=("agent:edges",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, ) # TODO use update_state once we have message ids