diff --git a/libs/checkpoint/README.md b/libs/checkpoint/README.md index d15480132..a78637471 100644 --- a/libs/checkpoint/README.md +++ b/libs/checkpoint/README.md @@ -10,21 +10,21 @@ Checkpoint is a snapshot of the graph state at a given point in time. Checkpoint ### Thread -Threads enable the checkpointing of multiple different runs, making them essential for multi-tenant chat applications and other scenarios where maintaining separate states is necessary. A thread is a unique ID assigned to a series of checkpoints saved by a checkpointer. When using a checkpointer, you must specify a `thread_id` or `thread_ts` when running the graph. +Threads enable the checkpointing of multiple different runs, making them essential for multi-tenant chat applications and other scenarios where maintaining separate states is necessary. A thread is a unique ID assigned to a series of checkpoints saved by a checkpointer. When using a checkpointer, you must specify a `thread_id` and optionally `checkpoint_id` when running the graph. - `thread_id` is simply the ID of a thread. This is always required -- `thread_ts` can optionally be passed. This identifier refers to a specific checkpoint within a thread. This can be used to kick of a run of a graph from some point halfway through a thread. +- `checkpoint_id` can optionally be passed. This identifier refers to a specific checkpoint within a thread. This can be used to kick of a run of a graph from some point halfway through a thread. You must pass these when invoking the graph as part of the configurable part of the config, e.g. ```python {"configurable": {"thread_id": "1"}} # valid config -{"configurable": {"thread_id": "1", "thread_ts": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}} # also valid config +{"configurable": {"thread_id": "1", "checkpoint_id": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}} # also valid config ``` ### Serde -`langgraph_checkpoint` also defines protocol for serialization/deserialization (serde) and provides an default implementation (`langgraph_checkpoint.serde.jsonplus.JsonPlusSerializer`) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more. +`langgraph_checkpoint` also defines protocol for serialization/deserialization (serde) and provides an default implementation (`langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer`) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more. ### Pending writes @@ -32,7 +32,7 @@ When a graph node fails mid-execution at a given superstep, LangGraph stores pen ## Interface -Each checkpointer should conform to `langgraph_checkpoint.BaseCheckpointSaver` interface and must implement the following methods: +Each checkpointer should conform to `langgraph.checkpoint.base.BaseCheckpointSaver` interface and must implement the following methods: - `.put` - Store a checkpoint with its configuration and metadata. - `.put_writes` - Store intermediate writes linked to a checkpoint (i.e. pending writes). diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 5ea3099d4..c553ddd7e 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -19,7 +19,7 @@ from typing import ( from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig from langgraph.checkpoint.base.id import uuid6 -from langgraph.checkpoint.serde.base import SerializerProtocol +from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( ChannelProtocol, @@ -172,10 +172,19 @@ CheckpointThreadId = ConfigurableFieldSpec( is_shared=True, ) -CheckpointThreadTs = ConfigurableFieldSpec( - id="thread_ts", +CheckpointNS = ConfigurableFieldSpec( + id="checkpoint_ns", + annotation=str, + name="Checkpoint NS", + description='Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', + default=None, + is_shared=True, +) + +CheckpointId = ConfigurableFieldSpec( + id="checkpoint_id", annotation=Optional[str], - name="Thread Timestamp", + name="Checkpoint ID", description="Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", default=None, is_shared=True, @@ -203,7 +212,7 @@ class BaseCheckpointSaver(ABC): *, serde: Optional[SerializerProtocol] = None, ) -> None: - self.serde = serde or self.serde + self.serde = maybe_add_typed_methods(serde or self.serde) @property def config_specs(self) -> list[ConfigurableFieldSpec]: @@ -212,7 +221,7 @@ class BaseCheckpointSaver(ABC): Returns: list[ConfigurableFieldSpec]: List of configuration field specs. """ - return [CheckpointThreadId, CheckpointThreadTs] + return [CheckpointThreadId, CheckpointNS, CheckpointId] def get(self, config: RunnableConfig) -> Optional[Checkpoint]: """Fetch a checkpoint using the given configuration. @@ -414,3 +423,10 @@ class EmptyChannelError(Exception): for the first time yet.""" pass + + +def get_checkpoint_id(config: RunnableConfig) -> Optional[str]: + """Get checkpoint ID in a backwards-compatible manner (fallback on thread_ts).""" + return config["configurable"].get( + "checkpoint_id", config["configurable"].get("thread_ts") + ) diff --git a/libs/checkpoint/langgraph/checkpoint/memory.py b/libs/checkpoint/langgraph/checkpoint/memory.py index 72b8c93db..db8da8ad0 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory.py +++ b/libs/checkpoint/langgraph/checkpoint/memory.py @@ -11,6 +11,7 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, SerializerProtocol, + get_checkpoint_id, ) @@ -44,7 +45,8 @@ class MemorySaver(BaseCheckpointSaver): asyncio.run(coro) # Output: 2 """ - storage: defaultdict[str, dict[str, tuple[bytes, bytes, Optional[str]]]] + # thread ID -> checkpoint NS -> checkpoint ID -> checkpoint mapping + storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]] def __init__( self, @@ -52,14 +54,14 @@ class MemorySaver(BaseCheckpointSaver): serde: Optional[SerializerProtocol] = None, ) -> None: super().__init__(serde=serde) - self.storage = defaultdict(dict) + self.storage = defaultdict(lambda: defaultdict(dict)) self.writes = defaultdict(list) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the in-memory storage. This method retrieves a checkpoint tuple from the in-memory storage based on the - provided config. If the config contains a "thread_ts" key, the checkpoint with + provided config. If the config contains a "checkpoint_id" key, the checkpoint with the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint for the given thread ID is retrieved. @@ -70,45 +72,54 @@ class MemorySaver(BaseCheckpointSaver): Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ thread_id = config["configurable"]["thread_id"] - if ts := config["configurable"].get("thread_ts"): - if saved := self.storage[thread_id].get(ts): - checkpoint, metadata, parent_ts = saved - writes = self.writes[(thread_id, ts)] + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + 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)] return CheckpointTuple( config=config, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), + checkpoint=self.serde.loads_typed(checkpoint), + metadata=self.serde.loads_typed(metadata), pending_writes=[ - (id, c, self.serde.loads(v)) for id, c, v in writes + (id, c, self.serde.loads_typed(v)) for id, c, v in writes ], parent_config={ "configurable": { "thread_id": thread_id, - "thread_ts": parent_ts, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if parent_ts + if parent_checkpoint_id else None, ) else: - if checkpoints := self.storage[thread_id]: - ts = max(checkpoints.keys()) - checkpoint, metadata, parent_ts = checkpoints[ts] - writes = self.writes[(thread_id, ts)] + 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)] return CheckpointTuple( - config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), + config={ + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, + } + }, + checkpoint=self.serde.loads_typed(checkpoint), + metadata=self.serde.loads_typed(metadata), pending_writes=[ - (id, c, self.serde.loads(v)) for id, c, v in writes + (id, c, self.serde.loads_typed(v)) for id, c, v in writes ], parent_config={ "configurable": { "thread_id": thread_id, - "thread_ts": parent_ts, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if parent_ts + if parent_checkpoint_id else None, ) @@ -135,16 +146,25 @@ class MemorySaver(BaseCheckpointSaver): Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage + checkpoint_ns = ( + config["configurable"].get("checkpoint_ns", "") if config else "" + ) for thread_id in thread_ids: - for ts, (checkpoint, metadata_b, parent_ts) in sorted( - self.storage[thread_id].items(), key=lambda x: x[0], reverse=True + for checkpoint_id, (checkpoint, metadata_b, parent_checkpoint_id) in sorted( + self.storage[thread_id][checkpoint_ns].items(), + key=lambda x: x[0], + reverse=True, ): - # filter by thread_ts - if before and ts >= before["configurable"]["thread_ts"]: + # filter by checkpoint ID + if ( + before + and (before_checkpoint_id := get_checkpoint_id(before)) + and checkpoint_id >= before_checkpoint_id + ): continue # filter by metadata - metadata = self.serde.loads(metadata_b) + metadata = self.serde.loads_typed(metadata_b) if filter and not all( query_value == metadata[query_key] for query_key, query_value in filter.items() @@ -158,16 +178,23 @@ class MemorySaver(BaseCheckpointSaver): limit -= 1 yield CheckpointTuple( - config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, - checkpoint=self.serde.loads(checkpoint), + config={ + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, + } + }, + checkpoint=self.serde.loads_typed(checkpoint), metadata=metadata, parent_config={ "configurable": { "thread_id": thread_id, - "thread_ts": parent_ts, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if parent_ts + if parent_checkpoint_id else None, ) @@ -190,19 +217,22 @@ class MemorySaver(BaseCheckpointSaver): Returns: RunnableConfig: The updated config containing the saved checkpoint's timestamp. """ - self.storage[config["configurable"]["thread_id"]].update( + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"]["checkpoint_ns"] + self.storage[thread_id][checkpoint_ns].update( { checkpoint["id"]: ( - self.serde.dumps(checkpoint), - self.serde.dumps(metadata), - config["configurable"].get("thread_ts"), # parent + self.serde.dumps_typed(checkpoint), + self.serde.dumps_typed(metadata), + config["configurable"].get("checkpoint_id"), # parent ) } ) return { "configurable": { - "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["id"], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint["id"], } } @@ -226,9 +256,11 @@ class MemorySaver(BaseCheckpointSaver): RunnableConfig: The updated config containing the saved writes' timestamp. """ thread_id = config["configurable"]["thread_id"] - ts = config["configurable"]["thread_ts"] - self.writes[(thread_id, ts)].extend( - [(task_id, c, self.serde.dumps(v)) for c, v in writes] + checkpoint_ns = config["configurable"]["checkpoint_ns"] + checkpoint_id = config["configurable"]["checkpoint_id"] + key = (thread_id, checkpoint_ns, checkpoint_id) + self.writes[key].extend( + [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes] ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: diff --git a/libs/checkpoint/langgraph/checkpoint/serde/base.py b/libs/checkpoint/langgraph/checkpoint/serde/base.py index 2fbb0ab71..5e4fba5be 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/base.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/base.py @@ -5,7 +5,9 @@ class SerializerProtocol(Protocol): """Protocol for serialization and deserialization of objects. - `dumps`: Serialize an object to bytes. + - `dumps_typed`: Serialize an object to a tuple (type, bytes). - `loads`: Deserialize an object from bytes. + - `loads_typed`: Deserialize an object from a tuple (type, bytes). Valid implementations include the `pickle`, `json` and `orjson` modules. """ @@ -13,5 +15,31 @@ class SerializerProtocol(Protocol): def dumps(self, obj: Any) -> bytes: ... + def dumps_typed(self, obj: Any) -> tuple[str, bytes]: + ... + def loads(self, data: bytes) -> Any: ... + + def loads_typed(self, data: tuple[str, bytes]) -> Any: + ... + + +class SerializerCompat(SerializerProtocol): + def __init__(self, serde: SerializerProtocol) -> None: + self.serde = serde + + def dumps_typed(self, obj: Any) -> tuple[str, bytes]: + return type(obj).__name__, self.serde.dumps(obj) + + def loads_typed(self, data: tuple[str, bytes]) -> Any: + return self.serde.loads(data[1]) + + +def maybe_add_typed_methods(serde: SerializerProtocol) -> SerializerProtocol: + """Wrap serde old serde implementations in a class with loads_typed and dumps_typed for backwards compatibility.""" + + if not hasattr(serde, "loads_typed") or not hasattr(serde, "dumps_typed"): + return SerializerCompat(serde) + + return serde diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index c76d90fe4..3d04a1f9a 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -99,5 +99,14 @@ class JsonPlusSerializer(SerializerProtocol): "utf-8", "ignore" ) + def dumps_typed(self, obj: Any) -> tuple[str, bytes]: + return "json", self.dumps(obj) + def loads(self, data: bytes) -> Any: return json.loads(data, object_hook=self._reviver) + + def loads_typed(self, data: tuple[str, bytes]) -> Any: + type_, data_ = data + if type_ != "json": + raise ValueError("JsonPlusSerializer can only deserialize `json` data") + return self.loads(data_) diff --git a/libs/checkpoint/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint/langgraph/checkpoint/sqlite/__init__.py index 48a34106b..78b601cb8 100644 --- a/libs/checkpoint/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/sqlite/__init__.py @@ -15,9 +15,11 @@ from langgraph.checkpoint.base import ( CheckpointTuple, EmptyChannelError, SerializerProtocol, + get_checkpoint_id, ) +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ChannelProtocol -from langgraph.checkpoint.sqlite.utils import JsonPlusSerializerCompat, search_where +from langgraph.checkpoint.sqlite.utils import search_where _AIO_ERROR_MSG = ( "The SqliteSaver does not support async methods. " @@ -61,11 +63,9 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): >>> graph.get_state(config) >>> result = graph.invoke(3, config) >>> graph.get_state(config) - StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '1', 'thread_ts': '2024-05-04T06:32:42.235444+00:00'}}, parent_config=None) + StateSnapshot(values=4, next=(), config={'configurable': {'thread_id': '1', 'checkpoint_id': '0c62ca34-ac19-445d-bbb0-5b4984975b2a'}}, parent_config=None) """ # noqa - serde = JsonPlusSerializerCompat() - conn: sqlite3.Connection is_setup: bool @@ -76,6 +76,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): serde: Optional[SerializerProtocol] = None, ) -> None: super().__init__(serde=serde) + self.jsonplus_serde = JsonPlusSerializer() self.conn = conn self.is_setup = False self.lock = threading.Lock() @@ -134,20 +135,24 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): PRAGMA journal_mode=WAL; CREATE TABLE IF NOT EXISTS checkpoints ( thread_id TEXT NOT NULL, - thread_ts TEXT NOT NULL, - parent_ts TEXT, + checkpoint_ns TEXT NOT NULL DEFAULT '', + checkpoint_id TEXT NOT NULL, + parent_checkpoint_id TEXT, + type TEXT, checkpoint BLOB, metadata BLOB, - PRIMARY KEY (thread_id, thread_ts) + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id) ); CREATE TABLE IF NOT EXISTS writes ( thread_id TEXT NOT NULL, - thread_ts TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL DEFAULT '', + checkpoint_id TEXT NOT NULL, task_id TEXT NOT NULL, idx INTEGER NOT NULL, channel TEXT NOT NULL, + type TEXT, value BLOB, - PRIMARY KEY (thread_id, thread_ts, task_id, idx) + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) ); """ ) @@ -180,7 +185,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): """Get a checkpoint tuple from the database. This method retrieves a checkpoint tuple from the SQLite database based on the - provided config. If the config contains a "thread_ts" key, the checkpoint with + provided config. If the config contains a "checkpoint_id" key, the checkpoint with the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint for the given thread ID is retrieved. @@ -203,63 +208,76 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): >>> config = { ... "configurable": { ... "thread_id": "1", - ... "thread_ts": "2024-05-04T06:32:42.235444+00:00", + ... "checkpoint_id": "2024-05-04T06:32:42.235444+00:00", ... } ... } >>> checkpoint_tuple = memory.get_tuple(config) >>> print(checkpoint_tuple) CheckpointTuple(...) """ # noqa + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") with self.cursor(transaction=False) as cur: # find the latest checkpoint for the thread_id - if config["configurable"].get("thread_ts"): + if checkpoint_id := get_checkpoint_id(config): cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", + "SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", ( str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), + checkpoint_ns, + checkpoint_id, ), ) else: cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1", - (str(config["configurable"]["thread_id"]),), + "SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1", + (str(config["configurable"]["thread_id"]), checkpoint_ns), ) # if a checkpoint is found, return it if value := cur.fetchone(): - if not config["configurable"].get("thread_ts"): + ( + thread_id, + checkpoint_id, + parent_checkpoint_id, + type, + checkpoint, + metadata, + ) = value + if not get_checkpoint_id(config): config = { "configurable": { - "thread_id": value[0], - "thread_ts": value[1], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, } } # find any pending writes cur.execute( - "SELECT task_id, channel, value FROM writes WHERE thread_id = ? AND thread_ts = ?", + "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", ( str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), + checkpoint_ns, + str(config["configurable"]["checkpoint_id"]), ), ) # deserialize the checkpoint and metadata return CheckpointTuple( config, - self.serde.loads(value[3]), - self.serde.loads(value[4]) if value[4] is not None else {}, + self.serde.loads_typed((type, checkpoint)), + self.jsonplus_serde.loads(metadata) if metadata is not None else {}, ( { "configurable": { - "thread_id": value[0], - "thread_ts": value[2], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if value[2] + if parent_checkpoint_id else None ), [ - (task_id, channel, self.serde.loads(value)) - for task_id, channel, value in cur + (task_id, channel, self.serde.loads_typed((type, value))) + for task_id, channel, type, value in cur ], ) @@ -295,33 +313,48 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): [CheckpointTuple(...), CheckpointTuple(...)] >>> config = {"configurable": {"thread_id": "1"}} - >>> before = {"configurable": {"thread_ts": "2024-05-04T06:32:42.235444+00:00"}} + >>> before = {"configurable": {"checkpoint_id": "2024-05-04T06:32:42.235444+00:00"}} >>> checkpoints = list(memory.list(config, before=before)) >>> print(checkpoints) [CheckpointTuple(...), ...] """ where, param_values = search_where(config, filter, before) - query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata + query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints {where} - ORDER BY thread_ts DESC""" + ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" with self.cursor(transaction=False) as cur: cur.execute(query, param_values) - for thread_id, thread_ts, parent_ts, value, metadata in cur: + for ( + thread_id, + checkpoint_ns, + checkpoint_id, + parent_checkpoint_id, + type, + checkpoint, + metadata, + ) in cur: yield CheckpointTuple( - {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, - self.serde.loads(value), - self.serde.loads(metadata) if metadata is not None else {}, + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, + } + }, + self.serde.loads_typed((type, checkpoint)), + self.jsonplus_serde.loads(metadata) if metadata is not None else {}, ( { "configurable": { "thread_id": thread_id, - "thread_ts": parent_ts, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if parent_ts + if parent_checkpoint_id else None ), ) @@ -354,23 +387,30 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): >>> checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "data": {"key": "value"}} >>> saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}) >>> print(saved_config) - {"configurable": {"thread_id": "1", "thread_ts": 2024-05-04T06:32:42.235444+00:00"}} + {"configurable": {"thread_id": "1", "checkpoint_id": 2024-05-04T06:32:42.235444+00:00"}} """ + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"]["checkpoint_ns"] + type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) + serialized_metadata = self.jsonplus_serde.dumps(metadata) with self.lock, self.cursor() as cur: cur.execute( - "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)", + "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), + checkpoint_ns, checkpoint["id"], - config["configurable"].get("thread_ts"), - self.serde.dumps(checkpoint), - self.serde.dumps(metadata), + config["configurable"].get("checkpoint_id"), + type_, + serialized_checkpoint, + serialized_metadata, ), ) return { "configurable": { - "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["id"], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint["id"], } } @@ -391,15 +431,16 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): """ with self.lock, self.cursor() as cur: cur.executemany( - "INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)", + "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ ( str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), + str(config["configurable"]["checkpoint_ns"]), + str(config["configurable"]["checkpoint_id"]), task_id, idx, channel, - self.serde.dumps(value), + *self.serde.dumps_typed(value), ) for idx, (channel, value) in enumerate(writes) ], @@ -463,7 +504,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): current_v = int(current.split(".")[0]) next_v = current_v + 1 try: - next_h = md5(self.serde.dumps(channel.checkpoint())).hexdigest() + next_h = md5(self.serde.dumps_typed(channel.checkpoint())[1]).hexdigest() except EmptyChannelError: next_h = "" return f"{next_v:032}.{next_h}" diff --git a/libs/checkpoint/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint/langgraph/checkpoint/sqlite/aio.py index 77b4486ee..69e07ad47 100644 --- a/libs/checkpoint/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint/langgraph/checkpoint/sqlite/aio.py @@ -23,8 +23,10 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, SerializerProtocol, + get_checkpoint_id, ) -from langgraph.checkpoint.sqlite.utils import JsonPlusSerializerCompat, search_where +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.checkpoint.sqlite.utils import search_where T = TypeVar("T", bound=callable) @@ -113,12 +115,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): ... saved_config = await saver.aput(config, checkpoint) ... print(saved_config) >>> asyncio.run(main()) - {"configurable": {"thread_id": "1", "thread_ts": "2023-05-03T10:00:00Z"}} + {"configurable": {"thread_id": "1", "checkpoint_id": "0c62ca34-ac19-445d-bbb0-5b4984975b2a"}} ``` """ - serde = JsonPlusSerializerCompat() - lock: asyncio.Lock is_setup: bool @@ -129,6 +129,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): serde: Optional[SerializerProtocol] = None, ): super().__init__(serde=serde) + self.jsonplus_serde = JsonPlusSerializer() self.conn = conn self.lock = asyncio.Lock() self.is_setup = False @@ -209,20 +210,24 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): PRAGMA journal_mode=WAL; CREATE TABLE IF NOT EXISTS checkpoints ( thread_id TEXT NOT NULL, - thread_ts TEXT NOT NULL, - parent_ts TEXT, + checkpoint_ns TEXT NOT NULL DEFAULT '', + checkpoint_id TEXT NOT NULL, + parent_checkpoint_id TEXT, + type TEXT, checkpoint BLOB, metadata BLOB, - PRIMARY KEY (thread_id, thread_ts) + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id) ); CREATE TABLE IF NOT EXISTS writes ( thread_id TEXT NOT NULL, - thread_ts TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL DEFAULT '', + checkpoint_id TEXT NOT NULL, task_id TEXT NOT NULL, idx INTEGER NOT NULL, channel TEXT NOT NULL, + type TEXT, value BLOB, - PRIMARY KEY (thread_id, thread_ts, task_id, idx) + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) ); """ ): @@ -234,7 +239,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): """Get a checkpoint tuple from the database asynchronously. This method retrieves a checkpoint tuple from the SQLite database based on the - provided config. If the config contains a "thread_ts" key, the checkpoint with + provided config. If the config contains a "checkpoint_id" key, the checkpoint with the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint for the given thread ID is retrieved. @@ -245,56 +250,69 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ await self.setup() + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") async with self.conn.cursor() as cur: # find the latest checkpoint for the thread_id - if config["configurable"].get("thread_ts"): + if checkpoint_id := get_checkpoint_id(config): await cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND thread_ts = ?", + "SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", ( str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), + checkpoint_ns, + checkpoint_id, ), ) else: await cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1", - (str(config["configurable"]["thread_id"]),), + "SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1", + (str(config["configurable"]["thread_id"]), checkpoint_ns), ) # if a checkpoint is found, return it if value := await cur.fetchone(): - if not config["configurable"].get("thread_ts"): + ( + thread_id, + checkpoint_id, + parent_checkpoint_id, + type, + checkpoint, + metadata, + ) = value + if not get_checkpoint_id(config): config = { "configurable": { - "thread_id": value[0], - "thread_ts": value[1], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, } } # find any pending writes await cur.execute( - "SELECT task_id, channel, value FROM writes WHERE thread_id = ? AND thread_ts = ?", + "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", ( str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), + checkpoint_ns, + str(config["configurable"]["checkpoint_id"]), ), ) # deserialize the checkpoint and metadata return CheckpointTuple( config, - self.serde.loads(value[3]), - self.serde.loads(value[4]) if value[4] is not None else {}, + self.serde.loads_typed((type, checkpoint)), + self.jsonplus_serde.loads(metadata) if metadata is not None else {}, ( { "configurable": { - "thread_id": value[0], - "thread_ts": value[2], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if value[2] + if parent_checkpoint_id else None ), [ - (task_id, channel, self.serde.loads(value)) - async for task_id, channel, value in cur + (task_id, channel, self.serde.loads_typed((type, value))) + async for task_id, channel, type, value in cur ], ) @@ -322,26 +340,41 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): """ await self.setup() where, param_values = search_where(config, filter, before) - query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata + query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints {where} - ORDER BY thread_ts DESC""" + ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" async with self.conn.execute(query, param_values) as cursor: - async for thread_id, thread_ts, parent_ts, value, metadata in cursor: + async for ( + thread_id, + checkpoint_ns, + checkpoint_id, + parent_checkpoint_id, + type, + checkpoint, + metadata, + ) in cursor: yield CheckpointTuple( - {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, - self.serde.loads(value), - self.serde.loads(metadata) if metadata is not None else {}, + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, + } + }, + self.serde.loads_typed((type, checkpoint)), + self.jsonplus_serde.loads(metadata) if metadata is not None else {}, ( { "configurable": { "thread_id": thread_id, - "thread_ts": parent_ts, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, } } - if parent_ts + if parent_checkpoint_id else None ), ) @@ -366,21 +399,28 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): RunnableConfig: The updated config containing the saved checkpoint's timestamp. """ await self.setup() + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"]["checkpoint_ns"] + type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) + serialized_metadata = self.jsonplus_serde.dumps(metadata) async with self.conn.execute( - "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, parent_ts, checkpoint, metadata) VALUES (?, ?, ?, ?, ?)", + "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), + checkpoint_ns, checkpoint["id"], - config["configurable"].get("thread_ts"), - self.serde.dumps(checkpoint), - self.serde.dumps(metadata), + config["configurable"].get("checkpoint_id"), + type_, + serialized_checkpoint, + serialized_metadata, ), ): await self.conn.commit() return { "configurable": { - "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["id"], + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint["id"], } } @@ -401,15 +441,16 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): """ await self.setup() async with self.conn.executemany( - "INSERT OR REPLACE INTO writes (thread_id, thread_ts, task_id, idx, channel, value) VALUES (?, ?, ?, ?, ?, ?)", + "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ ( str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), + str(config["configurable"]["checkpoint_ns"]), + str(config["configurable"]["checkpoint_id"]), task_id, idx, channel, - self.serde.dumps(value), + *self.serde.dumps_typed(value), ) for idx, (channel, value) in enumerate(writes) ], diff --git a/libs/checkpoint/langgraph/checkpoint/sqlite/utils.py b/libs/checkpoint/langgraph/checkpoint/sqlite/utils.py index f9334f42e..6e1baf5ae 100644 --- a/libs/checkpoint/langgraph/checkpoint/sqlite/utils.py +++ b/libs/checkpoint/langgraph/checkpoint/sqlite/utils.py @@ -1,38 +1,9 @@ import json -import pickle from typing import Any, Dict, Optional, Sequence, Tuple from langchain_core.runnables import RunnableConfig -from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer - - -class JsonPlusSerializerCompat(JsonPlusSerializer): - """A serializer that supports loading pickled checkpoints for backwards compatibility. - - This serializer extends the JsonPlusSerializer and adds support for loading pickled - checkpoints. If the input data starts with b"\x80" and ends with b".", it is treated - as a pickled checkpoint and loaded using pickle.loads(). Otherwise, the default - JsonPlusSerializer behavior is used. - - Examples: - >>> import pickle - >>> from langgraph.checkpoint.sqlite import JsonPlusSerializerCompat - >>> - >>> serializer = JsonPlusSerializerCompat() - >>> pickled_data = pickle.dumps({"key": "value"}) - >>> loaded_data = serializer.loads(pickled_data) - >>> print(loaded_data) # Output: {"key": "value"} - >>> - >>> json_data = '{"key": "value"}'.encode("utf-8") - >>> loaded_data = serializer.loads(json_data) - >>> print(loaded_data) # Output: {"key": "value"} - """ - - def loads(self, data: bytes) -> Any: - if data.startswith(b"\x80") and data.endswith(b"."): - return pickle.loads(data) - return super().loads(data) +from langgraph.checkpoint.base import get_checkpoint_id def _metadata_predicate( @@ -99,6 +70,9 @@ def search_where( if config is not None: wheres.append("thread_id = ?") param_values.append(config["configurable"]["thread_id"]) + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + wheres.append("checkpoint_ns = ?") + param_values.append(checkpoint_ns) # construct predicate for metadata filter if filter: @@ -108,7 +82,7 @@ def search_where( # construct predicate for `before` if before is not None: - wheres.append("thread_ts < ?") - param_values.append(before["configurable"]["thread_ts"]) + wheres.append("checkpoint_id < ?") + param_values.append(get_checkpoint_id(before)) return ("WHERE " + " AND ".join(wheres) if wheres else "", param_values) diff --git a/libs/checkpoint/tests/test_aiosqlite.py b/libs/checkpoint/tests/test_aiosqlite.py index 35f038d73..07188cc38 100644 --- a/libs/checkpoint/tests/test_aiosqlite.py +++ b/libs/checkpoint/tests/test_aiosqlite.py @@ -17,14 +17,31 @@ class TestAsyncSqliteSaver: # objects for test setup self.config_1: RunnableConfig = { - "configurable": {"thread_id": "thread-1", "thread_ts": "1"} + "configurable": { + "thread_id": "thread-1", + # for backwards compatibility testing + "thread_ts": "1", + "checkpoint_ns": "", + } } self.config_2: RunnableConfig = { - "configurable": {"thread_id": "thread-2", "thread_ts": "2"} + "configurable": { + "thread_id": "thread-2", + "checkpoint_id": "2", + "checkpoint_ns": "", + } + } + self.config_3: RunnableConfig = { + "configurable": { + "thread_id": "thread-2", + "checkpoint_id": "2-inner", + "checkpoint_ns": "inner", + } } self.chkpnt_1: Checkpoint = empty_checkpoint() self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) + self.chkpnt_3: Checkpoint = empty_checkpoint() self.metadata_1: CheckpointMetadata = { "source": "input", @@ -38,12 +55,14 @@ class TestAsyncSqliteSaver: "writes": {"foo": "bar"}, "score": None, } + self.metadata_3: CheckpointMetadata = {} async def test_asearch(self): # set up test # save checkpoints await self.sqlite_saver.aput(self.config_1, self.chkpnt_1, self.metadata_1) await self.sqlite_saver.aput(self.config_2, self.chkpnt_2, self.metadata_2) + await self.sqlite_saver.aput(self.config_3, self.chkpnt_3, self.metadata_3) # call method / assertions query_1: CheckpointMetadata = {"source": "input"} # search by 1 key @@ -70,11 +89,38 @@ class TestAsyncSqliteSaver: search_results_3 = [ c async for c in sqlite_saver.alist(None, filter=query_3) ] - assert len(search_results_3) == 2 + assert len(search_results_3) == 3 search_results_4 = [ c async for c in sqlite_saver.alist(None, filter=query_4) ] assert len(search_results_4) == 0 + # search by config (defaults to root graph checkpoints) + search_results_5 = [ + c + async for c in self.sqlite_saver.alist( + {"configurable": {"thread_id": "thread-2"}} + ) + ] + assert len(search_results_5) == 1 + assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" + + # search by config and checkpoint_ns + search_results_6 = [ + c + async for c in self.sqlite_saver.alist( + { + "configurable": { + "thread_id": "thread-2", + "checkpoint_ns": "inner", + } + } + ) + ] + assert len(search_results_6) == 1 + assert ( + search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" + ) + # TODO: test before and limit params diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index 254238c1f..7ec853d6a 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -99,14 +99,14 @@ def test_serde_jsonplus() -> None: serde = JsonPlusSerializer() - dumped = serde.dumps(to_serialize) + dumped = serde.dumps_typed(to_serialize) - assert ( - dumped - == b"""{"uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""" + assert dumped == ( + "json", + b"""{"uid": {"lc": 2, "type": "constructor", "id": ["uuid", "UUID"], "method": null, "args": ["00000000000000000000000000000001"], "kwargs": {}}, "time": {"lc": 2, "type": "constructor", "id": ["datetime", "datetime"], "method": "fromisoformat", "args": ["2024-04-19T23:04:57.051022+23:59"], "kwargs": {}}, "my_slotted_class": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclassWSlots"], "method": null, "args": [], "kwargs": {"foo": "bar", "bar": 2}}, "my_dataclass": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyDataclass"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_enum": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyEnum"], "method": null, "args": ["foo"], "kwargs": {}}, "my_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "my_funny_pydantic": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyFunnyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}, "person": {"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "Person"], "method": null, "args": [], "kwargs": {"name": "foo"}}, "a_bool": true, "a_none": null, "a_str": "foo", "a_str_nuc": "foo\\u0000", "a_str_uc": "foo \xe2\x9b\xb0\xef\xb8\x8f", "a_str_ucuc": "foo \xe2\x9b\xb0\xef\xb8\x8f\\u0000", "a_str_ucucuc": "foo \\\\u26f0\\\\ufe0f", "text": ["Hello", "Python", "Surrogate", "Example", "String", "With", "Surrogates", "Embedded", "In", "The", "Text", "\xe6\x94\xb6\xe8\x8a\xb1\xf0\x9f\x99\x84\xc2\xb7\xe5\x88\xb0"], "an_int": 1, "a_float": 1.1, "runnable_map": {"lc": 1, "type": "constructor", "id": ["langchain", "schema", "runnable", "RunnableParallel"], "kwargs": {"steps__": {}}, "name": "RunnableParallel<>", "graph": {"nodes": [{"id": 0, "type": "schema", "data": "Parallel<>Input"}, {"id": 1, "type": "schema", "data": "Parallel<>Output"}], "edges": []}}}""", ) - assert serde.loads(dumped) == { + assert serde.loads_typed(dumped) == { **to_serialize, "text": [v.encode("utf-8", "ignore").decode() for v in to_serialize["text"]], } diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 1622444c3..f955b6509 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -17,14 +17,31 @@ class TestMemorySaver: # objects for test setup self.config_1: RunnableConfig = { - "configurable": {"thread_id": "thread-1", "thread_ts": "1"} + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + # for backwards compatibility testing + "thread_ts": "1", + } } self.config_2: RunnableConfig = { - "configurable": {"thread_id": "thread-2", "thread_ts": "2"} + "configurable": { + "thread_id": "thread-2", + "checkpoint_ns": "", + "checkpoint_id": "2", + } + } + self.config_3: RunnableConfig = { + "configurable": { + "thread_id": "thread-2", + "checkpoint_id": "2-inner", + "checkpoint_ns": "inner", + } } self.chkpnt_1: Checkpoint = empty_checkpoint() self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) + self.chkpnt_3: Checkpoint = empty_checkpoint() self.metadata_1: CheckpointMetadata = { "source": "input", @@ -38,12 +55,14 @@ class TestMemorySaver: "writes": {"foo": "bar"}, "score": None, } + self.metadata_3: CheckpointMetadata = {} async def test_search(self): # set up test # save checkpoints self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1) self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2) + self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3) # call method / assertions query_1: CheckpointMetadata = {"source": "input"} # search by 1 key @@ -68,6 +87,22 @@ class TestMemorySaver: search_results_4 = list(self.memory_saver.list(None, filter=query_4)) assert len(search_results_4) == 0 + # search by config (defaults to root graph checkpoints) + search_results_5 = list( + self.memory_saver.list({"configurable": {"thread_id": "thread-2"}}) + ) + assert len(search_results_5) == 1 + assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" + + # search by config and checkpoint_ns + search_results_6 = list( + self.memory_saver.list( + {"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}} + ) + ) + assert len(search_results_6) == 1 + assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" + # TODO: test before and limit params async def test_asearch(self): diff --git a/libs/checkpoint/tests/test_sqlite.py b/libs/checkpoint/tests/test_sqlite.py index d920aeb5e..04af1a47e 100644 --- a/libs/checkpoint/tests/test_sqlite.py +++ b/libs/checkpoint/tests/test_sqlite.py @@ -18,14 +18,31 @@ class TestSqliteSaver: # objects for test setup self.config_1: RunnableConfig = { - "configurable": {"thread_id": "thread-1", "thread_ts": "1"} + "configurable": { + "thread_id": "thread-1", + # for backwards compatibility testing + "thread_ts": "1", + "checkpoint_ns": "", + } } self.config_2: RunnableConfig = { - "configurable": {"thread_id": "thread-2", "thread_ts": "2"} + "configurable": { + "thread_id": "thread-2", + "checkpoint_id": "2", + "checkpoint_ns": "", + } + } + self.config_3: RunnableConfig = { + "configurable": { + "thread_id": "thread-2", + "checkpoint_id": "2-inner", + "checkpoint_ns": "inner", + } } self.chkpnt_1: Checkpoint = empty_checkpoint() self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1) + self.chkpnt_3: Checkpoint = empty_checkpoint() self.metadata_1: CheckpointMetadata = { "source": "input", @@ -46,6 +63,7 @@ class TestSqliteSaver: # save checkpoints self.sqlite_saver.put(self.config_1, self.chkpnt_1, self.metadata_1) self.sqlite_saver.put(self.config_2, self.chkpnt_2, self.metadata_2) + self.sqlite_saver.put(self.config_3, self.chkpnt_3, self.metadata_3) # call method / assertions query_1: CheckpointMetadata = {"source": "input"} # search by 1 key @@ -65,16 +83,32 @@ class TestSqliteSaver: assert search_results_2[0].metadata == self.metadata_2 search_results_3 = list(self.sqlite_saver.list(None, filter=query_3)) - assert len(search_results_3) == 2 + assert len(search_results_3) == 3 search_results_4 = list(self.sqlite_saver.list(None, filter=query_4)) assert len(search_results_4) == 0 + # search by config (defaults to root graph checkpoints) + search_results_5 = list( + self.sqlite_saver.list({"configurable": {"thread_id": "thread-2"}}) + ) + assert len(search_results_5) == 1 + assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" + + # search by config and checkpoint_ns + search_results_6 = list( + self.sqlite_saver.list( + {"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}} + ) + ) + assert len(search_results_6) == 1 + assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" + # TODO: test before and limit params def test_search_where(self): # call method / assertions - expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND thread_ts < ?" + expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND checkpoint_id < ?" expected_param_values_1 = ["input", 2, "{}", 1, "1"] assert search_where(None, self.metadata_1, self.config_1) == ( expected_predicate_1, diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index f3aeb6a2e..1b6ae95d7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -21,6 +21,8 @@ TAG_HIDDEN = "langsmith:hidden" START = "__start__" END = "__end__" +CHECKPOINT_NAMESPACE_SEPARATOR = "|" + class Send: """A message or packet to send to a specific node in the graph. diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index f55520223..139b4ad0a 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -25,7 +25,13 @@ from langchain_core.runnables.graph import Node as DrawableNode from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import END, START, TAG_HIDDEN, Send +from langgraph.constants import ( + CHECKPOINT_NAMESPACE_SEPARATOR, + END, + START, + TAG_HIDDEN, + Send, +) from langgraph.errors import InvalidUpdateError from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode @@ -154,6 +160,11 @@ class Graph: *, metadata: Optional[dict[str, Any]] = None, ) -> None: + if isinstance(node, str) and CHECKPOINT_NAMESPACE_SEPARATOR in node: + raise ValueError( + f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names." + ) + if self.compiled: logger.warning( "Adding a node to a graph that has already been compiled. This will " diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 702fbbfb9..ffe35e8bc 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -30,7 +30,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import TAG_HIDDEN +from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR, TAG_HIDDEN from langgraph.errors import InvalidUpdateError from langgraph.graph.graph import ( END, @@ -311,6 +311,12 @@ class StateGraph(Graph): raise ValueError(f"Node `{node}` already present.") if node == END or node == START: raise ValueError(f"Node `{node}` is reserved.") + + if CHECKPOINT_NAMESPACE_SEPARATOR in node: + raise ValueError( + f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names." + ) + try: if isfunction(action) and ( hints := get_type_hints(action.__call__) or get_type_hints(action) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 47c7df81b..71df7964a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -519,7 +519,14 @@ class Pregel( checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() step = saved.metadata.get("step", -1) if saved else -1 # merge configurable fields with previous checkpoint config - checkpoint_config = config + checkpoint_config = { + **config, + "configurable": { + **config["configurable"], + # TODO: add proper support for updating nested subgraph state + "checkpoint_ns": "", + }, + } if saved: checkpoint_config = { "configurable": { @@ -681,7 +688,14 @@ class Pregel( step = saved.metadata.get("step", -2) + 1 if saved else -1 # merge configurable fields with previous checkpoint config - checkpoint_config = config + checkpoint_config = { + **config, + "configurable": { + **config["configurable"], + # TODO: add proper support for updating nested subgraph state + "checkpoint_ns": "", + }, + } if saved: checkpoint_config = { "configurable": { diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 17fde77bd..7a1a2abbb 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -33,6 +33,7 @@ from langgraph.checkpoint.base import ( create_checkpoint, ) from langgraph.constants import ( + CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, @@ -256,6 +257,7 @@ def prepare_next_tasks( checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, ) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]: + parent_ns = config.get("configurable", {}).get("checkpoint_ns", "") tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] # Consume pending packets for packet in checkpoint["pending_sends"]: @@ -275,7 +277,14 @@ def prepare_next_tasks( "langgraph_triggers": triggers, "langgraph_task_idx": len(tasks), } - task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) + checkpoint_ns = ( + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}" + if parent_ns + else packet.node + ) + task_id = str( + uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) + ) writes = deque() tasks.append( PregelExecutableTask( @@ -349,13 +358,18 @@ def prepare_next_tasks( "langgraph_triggers": triggers, "langgraph_task_idx": len(tasks), } - task_id = str(uuid5(UUID(checkpoint["id"]), json.dumps(metadata))) - if parent_thread_id := config.get("configurable", {}).get( - "thread_id" - ): - thread_id: Optional[str] = f"{parent_thread_id}-{name}" - else: - thread_id = None + checkpoint_ns = ( + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{name}" + if parent_ns + else name + ) + task_id = str( + uuid5( + UUID(checkpoint["id"]), + json.dumps((checkpoint_ns, metadata)), + ) + ) + writes = deque() tasks.append( PregelExecutableTask( @@ -389,8 +403,8 @@ def prepare_next_tasks( ), CONFIG_KEY_CHECKPOINTER: checkpointer, CONFIG_KEY_RESUMING: is_resuming, - "thread_id": thread_id, - "thread_ts": checkpoint["id"], + "checkpoint_id": checkpoint["id"], + "checkpoint_ns": checkpoint_ns, }, ), triggers, diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 70c1e69b7..c08699e2b 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -71,7 +71,7 @@ def map_debug_tasks( continue metadata = config["metadata"].copy() - metadata.pop("thread_ts", None) + metadata.pop("checkpoint_id", None) yield { "type": "task", @@ -97,7 +97,9 @@ def map_debug_task_results( continue metadata = config["metadata"].copy() - metadata.pop("thread_ts", None) + metadata.pop("checkpoint_id", None) + # TODO: make task IDs deterministic in tests and reuse task IDs for payload ID + metadata.pop("checkpoint_ns", None) yield { "type": "task_result", diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 678ff1217..37519cb25 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -143,7 +143,10 @@ class PregelLoop: **self.checkpoint_config, "configurable": { **self.checkpoint_config["configurable"], - "thread_ts": self.checkpoint["id"], + "checkpoint_ns": self.config["configurable"].get( + "checkpoint_ns", "" + ), + "checkpoint_id": self.checkpoint["id"], }, }, writes, @@ -315,8 +318,19 @@ class PregelLoop: # this is achieved by writing child checkpoints as progress is made # (so that error recovery / resuming from interrupt don't lose work) # but doing so always with an id equal to that of the parent checkpoint - id=self.config["configurable"]["thread_ts"] if self.is_nested else None, + id=self.config["configurable"]["checkpoint_id"] + if self.is_nested + else None, ) + self.checkpoint_config = { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "checkpoint_ns": self.config["configurable"].get( + "checkpoint_ns", "" + ), + }, + } # save it, without blocking # if there's a previous checkpoint save in progress, wait for it # ensuring checkpointers receive checkpoints in order @@ -331,7 +345,7 @@ class PregelLoop: **self.checkpoint_config, "configurable": { **self.checkpoint_config["configurable"], - "thread_ts": self.checkpoint["id"], + "checkpoint_id": self.checkpoint["id"], }, } # produce debug output diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index 09c9e52af..5e9607ad5 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -15,17 +15,17 @@ from langgraph.checkpoint.memory import MemorySaver class NoopSerializer(SerializerProtocol): - def loads(self, data: bytes) -> Any: - return data + def loads_typed(self, data: tuple[str, bytes]) -> Any: + return data[1] - def dumps(self, obj: Any) -> bytes: - return obj + def dumps_typed(self, obj: Any) -> tuple[str, bytes]: + return "type", obj class MemorySaverAssertImmutable(MemorySaver): serde = NoopSerializer() - storage_for_copies: defaultdict[str, dict[str, Checkpoint]] + storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]] def __init__( self, @@ -34,7 +34,7 @@ class MemorySaverAssertImmutable(MemorySaver): put_sleep: Optional[float] = None, ) -> None: super().__init__(serde=serde) - self.storage_for_copies = defaultdict(dict) + self.storage_for_copies = defaultdict(lambda: defaultdict(dict)) self.put_sleep = put_sleep def put( @@ -49,14 +49,17 @@ class MemorySaverAssertImmutable(MemorySaver): time.sleep(self.put_sleep) # assert checkpoint hasn't been modified since last written thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"]["checkpoint_ns"] if saved := super().get(config): assert ( - self.serde.loads(self.storage_for_copies[thread_id][saved["id"]]) + self.serde.loads_typed( + self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]] + ) == saved ) - self.storage_for_copies[thread_id][checkpoint["id"]] = self.serde.dumps( - copy_checkpoint(checkpoint) - ) + self.storage_for_copies[thread_id][checkpoint_ns][ + checkpoint["id"] + ] = self.serde.dumps_typed(copy_checkpoint(checkpoint)) # call super to write checkpoint return super().put(config, checkpoint, metadata) @@ -91,23 +94,24 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): """ configurable = config["configurable"].copy() - # remove thread_ts to make testing simpler - thread_ts = configurable.pop("thread_ts", None) - - self.storage[config["configurable"]["thread_id"]].update( + # remove checkpoint_id to make testing simpler + checkpoint_id = configurable.pop("checkpoint_id", None) + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"]["checkpoint_ns"] + self.storage[thread_id][checkpoint_ns].update( { checkpoint["id"]: ( - self.serde.dumps(checkpoint), + self.serde.dumps_typed(checkpoint), # merge configurable fields and metadata - self.serde.dumps({**configurable, **metadata}), - thread_ts, + self.serde.dumps_typed({**configurable, **metadata}), + checkpoint_id, ) } ) return { "configurable": { "thread_id": config["configurable"]["thread_id"], - "thread_ts": checkpoint["id"], + "checkpoint_id": checkpoint["id"], } } diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4427b4613..01d02a1f6 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -628,7 +628,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 6, "writes": 5}, @@ -641,7 +642,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 5, "writes": None}, @@ -654,7 +656,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": 4, "writes": 3}, @@ -667,7 +670,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 3, "writes": None}, @@ -680,7 +684,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": 2, "writes": 20}, @@ -693,7 +698,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 1, "writes": 4}, @@ -706,7 +712,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 0, "writes": None}, @@ -719,7 +726,8 @@ def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": -1, "writes": 2}, @@ -808,7 +816,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, @@ -821,7 +830,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, @@ -834,7 +844,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, @@ -847,7 +858,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, @@ -860,7 +872,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, @@ -873,7 +886,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 0, "writes": None}, @@ -886,7 +900,8 @@ def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": -1, "writes": 1}, @@ -1416,13 +1431,17 @@ 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)["id"] + assert ( + state.config["configurable"]["checkpoint_id"] == 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)["id"] + assert ( + state.config["configurable"]["checkpoint_id"] == 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) @@ -1462,8 +1481,8 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: } # sorted descending assert ( - thread_1_history[0].config["configurable"]["thread_ts"] - > thread_1_history[1].config["configurable"]["thread_ts"] + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] ) # cursor pagination cursored = list( @@ -1478,18 +1497,18 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: # can get each checkpoint using aget with config assert ( memory.get(thread_1_history[0].config)["id"] - == thread_1_history[0].config["configurable"]["thread_ts"] + == thread_1_history[0].config["configurable"]["checkpoint_id"] ) assert ( memory.get(thread_1_history[1].config)["id"] - == thread_1_history[1].config["configurable"]["thread_ts"] + == thread_1_history[1].config["configurable"]["checkpoint_id"] ) thread_1_next_config = app.update_state(thread_1_history[1].config, 10) # update creates a new checkpoint assert ( - thread_1_next_config["configurable"]["thread_ts"] - > thread_1_history[0].config["configurable"]["thread_ts"] + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] ) # update makes new checkpoint child of the previous one assert ( @@ -1957,7 +1976,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ) assert ( app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ - "thread_ts" + "checkpoint_id" ] is not None ) @@ -6537,7 +6556,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": {"my_key": ""}, @@ -6560,7 +6580,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { @@ -6607,7 +6628,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { @@ -6654,7 +6676,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { @@ -6701,7 +6724,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { @@ -7920,7 +7944,13 @@ def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -7928,23 +7958,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -7964,7 +8014,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -7980,7 +8031,8 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -7990,7 +8042,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8002,14 +8055,21 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8017,23 +8077,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8088,7 +8168,13 @@ def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8096,23 +8182,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8128,7 +8234,13 @@ def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8136,23 +8248,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8177,7 +8309,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8193,7 +8326,8 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -8203,7 +8337,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8215,14 +8350,21 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8230,23 +8372,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8273,7 +8435,13 @@ def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8281,23 +8449,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8319,7 +8507,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8331,14 +8520,21 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8346,23 +8542,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8384,7 +8600,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8400,7 +8617,8 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -8410,7 +8628,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8422,14 +8641,21 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8437,23 +8663,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8464,7 +8710,7 @@ def test_nested_graph_interrupts( ), ] - # test restarting from thread_ts + # test restarting from checkpoint_id config = {"configurable": {"thread_id": "6"}} app = graph.compile(checkpointer=checkpointer) assert app.invoke({"my_key": "my value"}, config, debug=True) == { @@ -8475,7 +8721,13 @@ def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8483,23 +8735,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8510,7 +8782,10 @@ def test_nested_graph_interrupts( ), ] child_state_history = [ - c for c in app.get_state_history({"configurable": {"thread_id": "6-inner"}}) + c + for c in app.get_state_history( + {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} + ) ] assert child_state_history == [ StateSnapshot( @@ -8518,8 +8793,9 @@ def test_nested_graph_interrupts( next=(), config={ "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8535,8 +8811,9 @@ def test_nested_graph_interrupts( created_at=AnyStr(), parent_config={ "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), } }, ), @@ -8547,10 +8824,10 @@ def test_nested_graph_interrupts( # check that child snapshot matches id of parent child_snapshot = child_state_history[0] assert ( - child_snapshot.config["configurable"]["thread_ts"] - == state_history[0].config["configurable"]["thread_ts"] + child_snapshot.config["configurable"]["checkpoint_id"] + == state_history[0].config["configurable"]["checkpoint_id"] ) - # check resuming from interrupt w/ thread_ts + # check resuming from interrupt w/ checkpoint_id interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] before_interrupt_config = before_interrupt_state_snapshot.config # going to get to interrupt again here, so the output is None @@ -8562,7 +8839,13 @@ def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8570,13 +8853,23 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8584,23 +8877,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -8622,7 +8935,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8638,7 +8952,8 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -8648,7 +8963,8 @@ def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -8660,14 +8976,21 @@ def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8675,13 +8998,23 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -8689,23 +9022,43 @@ def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 83fe0951f..883dc2c40 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -686,7 +686,6 @@ async def test_invoke_two_processes_in_out_interrupt( add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel( nodes={"one": one, "two": two}, channels={ @@ -750,7 +749,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 6, "writes": 5}, @@ -763,7 +763,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 5, "writes": None}, @@ -776,7 +777,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": 4, "writes": 3}, @@ -789,7 +791,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 3, "writes": None}, @@ -802,7 +805,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": 2, "writes": 20}, @@ -815,7 +819,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 1, "writes": 4}, @@ -828,7 +833,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 0, "writes": None}, @@ -841,7 +847,8 @@ async def test_invoke_two_processes_in_out_interrupt( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": -1, "writes": 2}, @@ -908,7 +915,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, @@ -921,7 +929,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, @@ -934,7 +943,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, @@ -947,7 +957,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, @@ -960,7 +971,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, @@ -973,7 +985,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "loop", "step": 0, "writes": None}, @@ -986,7 +999,8 @@ async def test_fork_always_re_runs_nodes( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={"source": "input", "step": -1, "writes": 1}, @@ -1508,7 +1522,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: assert state is not None assert state.values.get("total") == 2 assert ( - state.config["configurable"]["thread_ts"] + state.config["configurable"]["checkpoint_id"] == (await memory.aget(thread_1))["id"] ) # total is now 2, so output is 2+3=5 @@ -1517,7 +1531,7 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: assert state is not None assert state.values.get("total") == 7 assert ( - state.config["configurable"]["thread_ts"] + state.config["configurable"]["checkpoint_id"] == (await memory.aget(thread_1))["id"] ) # total is now 2+5=7, so output would be 7+4=11, but raises ValueError @@ -1559,8 +1573,8 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: } # sorted descending assert ( - thread_1_history[0].config["configurable"]["thread_ts"] - > thread_1_history[1].config["configurable"]["thread_ts"] + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] ) # cursor pagination cursored = [ @@ -1578,16 +1592,16 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: # can get each checkpoint using aget with config assert (await memory.aget(thread_1_history[0].config))[ "id" - ] == thread_1_history[0].config["configurable"]["thread_ts"] + ] == thread_1_history[0].config["configurable"]["checkpoint_id"] assert (await memory.aget(thread_1_history[1].config))[ "id" - ] == thread_1_history[1].config["configurable"]["thread_ts"] + ] == thread_1_history[1].config["configurable"]["checkpoint_id"] thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) # update creates a new checkpoint assert ( - thread_1_next_config["configurable"]["thread_ts"] - > thread_1_history[0].config["configurable"]["thread_ts"] + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] ) # 1 more checkpoint in history assert len([c async for c in app.aget_state_history(thread_1)]) == 8 @@ -5161,7 +5175,11 @@ async def test_branch_then() -> None: "metadata": {"thread_id": "10"}, "callbacks": None, "recursion_limit": 25, - "configurable": {"thread_id": "10", "thread_ts": AnyStr()}, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, }, "values": {"my_key": ""}, "metadata": { @@ -5183,14 +5201,19 @@ async def test_branch_then() -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { "my_key": "value", "market": "DE", }, - "metadata": {"source": "loop", "step": 0, "writes": None}, + "metadata": { + "source": "loop", + "step": 0, + "writes": None, + }, }, }, { @@ -5226,10 +5249,14 @@ async def test_branch_then() -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, - "values": {"my_key": "value prepared", "market": "DE"}, + "values": { + "my_key": "value prepared", + "market": "DE", + }, "metadata": { "source": "loop", "step": 1, @@ -5270,7 +5297,8 @@ async def test_branch_then() -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { @@ -5317,7 +5345,8 @@ async def test_branch_then() -> None: "recursion_limit": 25, "configurable": { "thread_id": "10", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, }, "values": { @@ -6319,7 +6348,13 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6327,23 +6362,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6363,7 +6418,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6379,7 +6435,8 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -6389,7 +6446,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6401,14 +6459,21 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "1", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6416,23 +6481,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6496,7 +6581,13 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6504,23 +6595,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6536,7 +6647,13 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6544,23 +6661,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6585,7 +6722,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6601,7 +6739,8 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -6611,7 +6750,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6623,14 +6763,21 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "4", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6638,23 +6785,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "4", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6685,7 +6852,13 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6693,23 +6866,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6731,7 +6924,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6743,14 +6937,21 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6758,23 +6959,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6796,7 +7017,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6812,7 +7034,8 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -6822,7 +7045,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6834,14 +7058,21 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "5", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6849,23 +7080,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "5", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "5", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6876,7 +7127,7 @@ async def test_nested_graph_interrupts( ), ] - # test restarting from thread_ts + # test restarting from checkpoint_id config = {"configurable": {"thread_id": "6"}} app = graph.compile(checkpointer=checkpointer) await app.ainvoke({"my_key": "my value"}, config, debug=True) @@ -6886,7 +7137,13 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6894,23 +7151,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -6924,7 +7201,7 @@ async def test_nested_graph_interrupts( child_state_history = [ c async for c in app.aget_state_history( - {"configurable": {"thread_id": "6-inner"}} + {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} ) ] assert child_state_history == [ @@ -6933,8 +7210,9 @@ async def test_nested_graph_interrupts( next=(), config={ "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -6950,20 +7228,23 @@ async def test_nested_graph_interrupts( created_at=AnyStr(), parent_config={ "configurable": { - "thread_id": "6-inner", - "thread_ts": AnyStr(), + "thread_id": "6", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), } }, ), + # there should be a single child checkpoint because we only keep + # one child checkpoint per parent checkpoint (in which child ran) ] # check that child snapshot matches id of parent child_snapshot = child_state_history[0] assert ( - child_snapshot.config["configurable"]["thread_ts"] - == state_history[0].config["configurable"]["thread_ts"] + child_snapshot.config["configurable"]["checkpoint_id"] + == state_history[0].config["configurable"]["checkpoint_id"] ) - # check resuming from interrupt w/ thread_ts + # check resuming from interrupt w/ checkpoint_id interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] before_interrupt_config = before_interrupt_state_snapshot.config # going to get to interrupt again here @@ -6975,7 +7256,13 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6983,13 +7270,23 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -6997,23 +7294,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"}, @@ -7035,7 +7352,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -7051,7 +7369,8 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), @@ -7061,7 +7380,8 @@ async def test_nested_graph_interrupts( config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, metadata={ @@ -7073,14 +7393,21 @@ async def test_nested_graph_interrupts( parent_config={ "configurable": { "thread_id": "6", - "thread_ts": AnyStr(), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -7088,13 +7415,23 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, @@ -7102,23 +7439,43 @@ async def test_nested_graph_interrupts( }, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={"my_key": "my value"}, next=("outer_1",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={"source": "loop", "writes": None, "step": 0}, created_at=AnyStr(), parent_config={ - "configurable": {"thread_id": "6", "thread_ts": AnyStr()} + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, ), StateSnapshot( values={}, next=("__start__",), - config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}}, + config={ + "configurable": { + "thread_id": "6", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, metadata={ "source": "input", "writes": {"my_key": "my value"},