From 52b2c755d728d262061ebffd63adf0e7398cb62a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 May 2024 16:56:25 -0700 Subject: [PATCH] Use uuid6 as the id for checkpoints - this avoids conflicts if multiple processes creating checkpoints in same thread at same time - uuid6 with a monotically increasing clock_seq is sortable by creation time at ms precision (plus clock_seq for ties, plus 48 bits of randomness for further ties) --- langgraph/channels/base.py | 7 +++---- langgraph/checkpoint/aiosqlite.py | 4 ++-- langgraph/checkpoint/base.py | 6 ++++++ langgraph/checkpoint/memory.py | 4 ++-- langgraph/checkpoint/sqlite.py | 4 ++-- langgraph/pregel/__init__.py | 28 +++++++++++++++------------- poetry.lock | 13 ++++++++++++- pyproject.toml | 1 + tests/checkpoint/test_aiosqlite.py | 19 ++++--------------- tests/checkpoint/test_memory.py | 19 ++++--------------- tests/checkpoint/test_sqlite.py | 19 ++++--------------- tests/memory_assert.py | 4 ++-- tests/test_pregel.py | 14 +++++++------- tests/test_pregel_async.py | 14 +++++++------- 14 files changed, 71 insertions(+), 85 deletions(-) diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index ae6869b41..4384cc273 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -13,6 +13,7 @@ from typing import ( ) from typing_extensions import Self +from uuid6 import uuid6 from langgraph.checkpoint.base import Checkpoint from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -110,13 +111,10 @@ async def AsyncChannelsManager( def create_checkpoint( - checkpoint: Checkpoint, channels: Mapping[str, BaseChannel] + checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], step: int ) -> Checkpoint: """Create a checkpoint for the given channels.""" ts = datetime.now(timezone.utc).isoformat() - assert ( - ts > checkpoint["ts"] - ), f"Timestamps must be monotonically increasing, got {ts} <= {checkpoint['ts']}" values: dict[str, Any] = {} for k, v in channels.items(): try: @@ -126,6 +124,7 @@ def create_checkpoint( return Checkpoint( v=1, ts=ts, + id=str(uuid6(clock_seq=step)), channel_values=values, channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index b82cdfad8..91ac10c66 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -319,7 +319,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), - checkpoint["ts"], + checkpoint["id"], config["configurable"].get("thread_ts"), self.serde.dumps(checkpoint), self.serde.dumps(metadata), @@ -329,6 +329,6 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): return { "configurable": { "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 5a9dc99ce..1823e1eac 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -12,6 +12,7 @@ from typing import ( ) from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig +from uuid6 import uuid6 from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer @@ -48,6 +49,9 @@ class Checkpoint(TypedDict): v: int """The version of the checkpoint format. Currently 1.""" + id: str + """The ID of the checkpoint. This is both unique and monotonically + increasing, so can be used for sorting checkpoints from first to last.""" ts: str """The timestamp of the checkpoint in ISO 8601 format.""" channel_values: dict[str, Any] @@ -77,6 +81,7 @@ def _seen_dict(): def empty_checkpoint() -> Checkpoint: return Checkpoint( v=1, + id=str(uuid6(clock_seq=-2)), ts=datetime.now(timezone.utc).isoformat(), channel_values={}, channel_versions=defaultdict(int), @@ -88,6 +93,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: return Checkpoint( v=checkpoint["v"], ts=checkpoint["ts"], + id=checkpoint["id"], channel_values=checkpoint["channel_values"].copy(), channel_versions=defaultdict(int, checkpoint["channel_versions"]), versions_seen=defaultdict( diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index f8c49ce78..2f592cec5 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -192,7 +192,7 @@ class MemorySaver(BaseCheckpointSaver): """ self.storage[config["configurable"]["thread_id"]].update( { - checkpoint["ts"]: ( + checkpoint["id"]: ( self.serde.dumps(checkpoint), self.serde.dumps(metadata), ) @@ -201,7 +201,7 @@ class MemorySaver(BaseCheckpointSaver): return { "configurable": { "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index e5835f17b..83f1568a9 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -426,7 +426,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), - checkpoint["ts"], + checkpoint["id"], config["configurable"].get("thread_ts"), self.serde.dumps(checkpoint), self.serde.dumps(metadata), @@ -435,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): return { "configurable": { "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e6fffd3ce..210d75028 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -538,12 +538,13 @@ class Pregel( ) # apply to checkpoint and save _apply_writes(checkpoint, channels, task.writes) + step = saved.metadata.get("step", -2) + 1 if saved else -1 return self.checkpointer.put( saved.config if saved else config, - create_checkpoint(checkpoint, channels), + create_checkpoint(checkpoint, channels, step), { "source": "update", - "step": saved.metadata.get("step", 0) + 1 if saved else 0, + "step": step, "writes": {as_node: values}, }, ) @@ -612,12 +613,13 @@ class Pregel( ) # apply to checkpoint and save _apply_writes(checkpoint, channels, task.writes) + step = saved.metadata.get("step", -2) + 1 if saved else -1 return await self.checkpointer.aput( saved.config if saved else config, - create_checkpoint(checkpoint, channels), + create_checkpoint(checkpoint, channels, step), { "source": "update", - "step": saved.metadata.get("step", 0) + 1 if saved else 0, + "step": step, "writes": {as_node: values}, }, ) @@ -741,7 +743,7 @@ class Pregel( _apply_writes(checkpoint, channels, input_writes) # save input checkpoint if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels) + checkpoint = create_checkpoint(checkpoint, channels, start) bg.append( executor.submit( self.checkpointer.put, @@ -753,7 +755,7 @@ class Pregel( checkpoint_config = { "configurable": { **checkpoint_config["configurable"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } # increment start to 0 @@ -874,7 +876,7 @@ class Pregel( # save end of step checkpoint if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels) + checkpoint = create_checkpoint(checkpoint, channels, step) bg.append( executor.submit( self.checkpointer.put, @@ -898,7 +900,7 @@ class Pregel( checkpoint_config = { "configurable": { **checkpoint_config["configurable"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } # yield debug checkpoint @@ -1031,7 +1033,7 @@ class Pregel( _apply_writes(checkpoint, channels, input_writes) # save input checkpoint if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels) + checkpoint = create_checkpoint(checkpoint, channels, start) bg.append( asyncio.create_task( self.checkpointer.aput( @@ -1044,7 +1046,7 @@ class Pregel( checkpoint_config = { "configurable": { **checkpoint_config["configurable"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } # increment start to 0 @@ -1175,7 +1177,7 @@ class Pregel( # save end of step checkpoint if self.checkpointer is not None: - checkpoint = create_checkpoint(checkpoint, channels) + checkpoint = create_checkpoint(checkpoint, channels, step) bg.append( asyncio.create_task( self.checkpointer.aput( @@ -1200,7 +1202,7 @@ class Pregel( checkpoint_config = { "configurable": { **checkpoint_config["configurable"], - "thread_ts": checkpoint["ts"], + "thread_ts": checkpoint["id"], } } # yield debug checkpoint @@ -1414,7 +1416,7 @@ def _local_read( fresh: bool = False, ) -> Union[dict[str, Any], Any]: if fresh: - checkpoint = create_checkpoint(checkpoint, channels) + checkpoint = create_checkpoint(checkpoint, channels, -1) with ChannelsManager(channels, checkpoint) as channels: _apply_writes(copy_checkpoint(checkpoint), channels, writes) return read_channels(channels, select) diff --git a/poetry.lock b/poetry.lock index 8420bf216..6173f80be 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3868,6 +3868,17 @@ h2 = ["h2 (>=4,<5)"] socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] zstd = ["zstandard (>=0.18.0)"] +[[package]] +name = "uuid6" +version = "2024.1.12" +description = "New time-based UUID formats which are suited for use as a database key" +optional = false +python-versions = ">=3.8" +files = [ + {file = "uuid6-2024.1.12-py3-none-any.whl", hash = "sha256:8150093c8d05a331bc0535bc5ef6cf57ac6eceb2404fd319bc10caee2e02c065"}, + {file = "uuid6-2024.1.12.tar.gz", hash = "sha256:ed0afb3a973057575f9883201baefe402787ca5e11e1d24e377190f0c43f1993"}, +] + [[package]] name = "watchdog" version = "4.0.0" @@ -4094,4 +4105,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "d42a7ee4e0079a588710997ce02f3717820706d93922990537d0e4ab78a78beb" +content-hash = "a2a48cce5d31d0c3e09a9b3b042586c569c3efbe881de7a9c4ce9ae2b4d63499" diff --git a/pyproject.toml b/pyproject.toml index 491f9a2b2..de41a6d8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ repository = "https://www.github.com/langchain-ai/langgraph" [tool.poetry.dependencies] python = ">=3.9.0,<4.0" langchain-core = "^0.1.52" +uuid6 = "^2024.1.12" [tool.poetry.group.test.dependencies] diff --git a/tests/checkpoint/test_aiosqlite.py b/tests/checkpoint/test_aiosqlite.py index ef06efdcd..fe05296ea 100644 --- a/tests/checkpoint/test_aiosqlite.py +++ b/tests/checkpoint/test_aiosqlite.py @@ -1,8 +1,9 @@ import pytest from langchain_core.runnables import RunnableConfig +from langgraph.channels.base import create_checkpoint from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver -from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata +from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint class TestAsyncSqliteSaver: @@ -18,20 +19,8 @@ class TestAsyncSqliteSaver: "configurable": {"thread_id": "thread-2", "thread_ts": "2"} } - self.chkpnt_1: Checkpoint = { - "v": 1, - "ts": "1", - "channel_values": {}, - "channel_versions": {}, - "versions_seen": {}, - } - self.chkpnt_2: Checkpoint = { - "v": 2, - "ts": "2", - "channel_values": {}, - "channel_versions": {}, - "versions_seen": {}, - } + self.chkpnt_1: Checkpoint = empty_checkpoint() + self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) self.metadata_1: CheckpointMetadata = { "source": "input", diff --git a/tests/checkpoint/test_memory.py b/tests/checkpoint/test_memory.py index 2a4f26e34..32de877a4 100644 --- a/tests/checkpoint/test_memory.py +++ b/tests/checkpoint/test_memory.py @@ -1,7 +1,8 @@ import pytest from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata +from langgraph.channels.base import create_checkpoint +from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint from langgraph.checkpoint.memory import MemorySaver @@ -18,20 +19,8 @@ class TestMemorySaver: "configurable": {"thread_id": "thread-2", "thread_ts": "2"} } - self.chkpnt_1: Checkpoint = { - "v": 1, - "ts": "1", - "channel_values": {}, - "channel_versions": {}, - "versions_seen": {}, - } - self.chkpnt_2: Checkpoint = { - "v": 2, - "ts": "2", - "channel_values": {}, - "channel_versions": {}, - "versions_seen": {}, - } + self.chkpnt_1: Checkpoint = empty_checkpoint() + self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) self.metadata_1: CheckpointMetadata = { "source": "input", diff --git a/tests/checkpoint/test_sqlite.py b/tests/checkpoint/test_sqlite.py index c64b5ceea..741a908e1 100644 --- a/tests/checkpoint/test_sqlite.py +++ b/tests/checkpoint/test_sqlite.py @@ -1,7 +1,8 @@ import pytest from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata +from langgraph.channels.base import create_checkpoint +from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, empty_checkpoint from langgraph.checkpoint.sqlite import SqliteSaver, _metadata_predicate, search_where @@ -18,20 +19,8 @@ class TestSqliteSaver: "configurable": {"thread_id": "thread-2", "thread_ts": "2"} } - self.chkpnt_1: Checkpoint = { - "v": 1, - "ts": "1", - "channel_values": {}, - "channel_versions": {}, - "versions_seen": {}, - } - self.chkpnt_2: Checkpoint = { - "v": 2, - "ts": "2", - "channel_values": {}, - "channel_versions": {}, - "versions_seen": {}, - } + self.chkpnt_1: Checkpoint = empty_checkpoint() + self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) self.metadata_1: CheckpointMetadata = { "source": "input", diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 2429cf63d..99a438538 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -40,8 +40,8 @@ class MemorySaverAssertImmutable(MemorySaver): # 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["ts"]] == saved - self.storage_for_copies[thread_id][checkpoint["ts"]] = copy_checkpoint( + assert self.storage_for_copies[thread_id][saved["id"]] == saved + self.storage_for_copies[thread_id][checkpoint["id"]] = copy_checkpoint( checkpoint ) # call super to write checkpoint diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 787946cab..b98000b89 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -818,13 +818,13 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: assert state is not None assert state.values.get("total") == 2 assert state.next == () - assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["ts"] + assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["id"] # 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"] + assert state.config["configurable"]["thread_ts"] == memory.get(thread_1)["id"] # 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) @@ -879,11 +879,11 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: assert thread_1_history[-2].values["total"] == 2 # can get each checkpoint using aget with config assert ( - memory.get(thread_1_history[0].config)["ts"] + memory.get(thread_1_history[0].config)["id"] == thread_1_history[0].config["configurable"]["thread_ts"] ) assert ( - memory.get(thread_1_history[1].config)["ts"] + memory.get(thread_1_history[1].config)["id"] == thread_1_history[1].config["configurable"]["thread_ts"] ) @@ -4629,7 +4629,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: config=uconfig, metadata={ "source": "update", - "step": 0, + "step": -1, "writes": {START: {"my_key": "key", "market": "DE"}}, }, ) @@ -4645,7 +4645,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: config=tool_two.checkpointer.get_tuple(thread3).config, metadata={ "source": "loop", - "step": 1, + "step": 0, "writes": {"prepare": {"my_key": " prepared"}}, }, parent_config=uconfig, @@ -4661,7 +4661,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: config=tool_two.checkpointer.get_tuple(thread3).config, metadata={ "source": "loop", - "step": 3, + "step": 2, "writes": {"finish": {"my_key": " finished"}}, }, parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index f6f0bb50b..40892621f 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -801,7 +801,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: assert state.values.get("total") == 2 assert ( state.config["configurable"]["thread_ts"] - == (await memory.aget(thread_1))["ts"] + == (await memory.aget(thread_1))["id"] ) # total is now 2, so output is 2+3=5 assert await app.ainvoke(3, thread_1) == 5 @@ -810,7 +810,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: assert state.values.get("total") == 7 assert ( state.config["configurable"]["thread_ts"] - == (await memory.aget(thread_1))["ts"] + == (await memory.aget(thread_1))["id"] ) # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): @@ -869,10 +869,10 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: assert thread_1_history[-2].values["total"] == 2 # can get each checkpoint using aget with config assert (await memory.aget(thread_1_history[0].config))[ - "ts" + "id" ] == thread_1_history[0].config["configurable"]["thread_ts"] assert (await memory.aget(thread_1_history[1].config))[ - "ts" + "id" ] == thread_1_history[1].config["configurable"]["thread_ts"] thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) @@ -3907,7 +3907,7 @@ async def test_branch_then() -> None: config=uconfig, metadata={ "source": "update", - "step": 0, + "step": -1, "writes": {START: {"my_key": "key", "market": "DE"}}, }, ) @@ -3923,7 +3923,7 @@ async def test_branch_then() -> None: config=(await tool_two.checkpointer.aget_tuple(thread3)).config, metadata={ "source": "loop", - "step": 1, + "step": 0, "writes": {"prepare": {"my_key": " prepared"}}, }, parent_config=uconfig, @@ -3939,7 +3939,7 @@ async def test_branch_then() -> None: config=(await tool_two.checkpointer.aget_tuple(thread3)).config, metadata={ "source": "loop", - "step": 3, + "step": 2, "writes": {"finish": {"my_key": " finished"}}, }, parent_config=[