Add history tracking to in-memory, sqlite and aiosqlite checkpointers

- Add Pregel.get_state_history and .aget_state_history methods to get history iterator
- Update checkpointer base class with new list and get_tuple methods
- Rewrite in-memory checkpointer class to track history
- Rewrite sqlite and aiosqlite checkpointers to track history
- Add new tests for history tracking

langgraph/checkpoint/sqlite.py

@@ -1,16 +1,16 @@
 import pickle
 import sqlite3
-from contextlib import contextmanager
-from typing import Optional
+from contextlib import AbstractContextManager, contextmanager
+from types import TracebackType
+from typing import AsyncIterator, Iterator, Optional, Self
 
 from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.utils import ConfigurableFieldSpec
 
-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
+from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple
 
 
-class SqliteSaver(BaseCheckpointSaver):
+class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
     conn: sqlite3.Connection
 
     is_setup: bool = Field(False, init=False, repr=False)
@@ -22,18 +22,16 @@ class SqliteSaver(BaseCheckpointSaver):
     def from_conn_string(cls, conn_string: str) -> "SqliteSaver":
         return SqliteSaver(conn=sqlite3.connect(conn_string))
 
-    @property
-    def config_specs(self) -> list[ConfigurableFieldSpec]:
-        return [
-            ConfigurableFieldSpec(
-                id="thread_id",
-                annotation=str,
-                name="Thread ID",
-                description=None,
-                default="",
-                is_shared=True,
-            ),
-        ]
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        __exc_type: type[BaseException] | None,
+        __exc_value: BaseException | None,
+        __traceback: TracebackType | None,
+    ) -> bool | None:
+        return self.conn.close()
 
     def setup(self) -> None:
         if self.is_setup:
@@ -42,8 +40,10 @@ class SqliteSaver(BaseCheckpointSaver):
         self.conn.executescript(
             """
             CREATE TABLE IF NOT EXISTS checkpoints (
-                thread_id TEXT PRIMARY KEY,
-                checkpoint BLOB
+                thread_id TEXT NOT NULL,
+                thread_ts TEXT NOT NULL,
+                checkpoint BLOB,
+                PRIMARY KEY (thread_id, thread_ts)
             );
             """
         )
@@ -61,27 +61,72 @@ class SqliteSaver(BaseCheckpointSaver):
                 self.conn.commit()
             cur.close()
 
-    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        with self.cursor(transaction=False) as cur:
+            if config["configurable"].get("thread_ts"):
+                cur.execute(
+                    "SELECT checkpoint FROM checkpoints WHERE thread_id = ? AND thread_ts = ?",
+                    (
+                        config["configurable"]["thread_id"],
+                        config["configurable"]["thread_ts"],
+                    ),
+                )
+                if value := cur.fetchone():
+                    return CheckpointTuple(config, pickle.loads(value[0]))
+            else:
+                cur.execute(
+                    "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC LIMIT 1",
+                    (config["configurable"]["thread_id"],),
+                )
+                if value := cur.fetchone():
+                    return CheckpointTuple(
+                        {
+                            "configurable": {
+                                "thread_id": value[0],
+                                "thread_ts": value[1],
+                            }
+                        },
+                        pickle.loads(value[2]),
+                    )
+
+    def list(self, config: RunnableConfig) -> Iterator[CheckpointTuple]:
         with self.cursor(transaction=False) as cur:
             cur.execute(
-                "SELECT checkpoint FROM checkpoints WHERE thread_id = ?",
+                "SELECT thread_id, thread_ts, checkpoint FROM checkpoints WHERE thread_id = ? ORDER BY thread_ts DESC",
                 (config["configurable"]["thread_id"],),
             )
-            if value := cur.fetchone():
-                return pickle.loads(value[0])
+            for thread_id, thread_ts, value in cur:
+                yield CheckpointTuple(
+                    {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
+                    pickle.loads(value),
+                )
 
-    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
         with self.cursor() as cur:
             cur.execute(
-                "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)",
+                "INSERT OR REPLACE INTO checkpoints (thread_id, thread_ts, checkpoint) VALUES (?, ?, ?)",
                 (
                     config["configurable"]["thread_id"],
+                    checkpoint["ts"],
                     pickle.dumps(checkpoint),
                 ),
             )
+        return {
+            "configurable": {
+                "thread_id": config["configurable"]["thread_id"],
+                "thread_ts": checkpoint["ts"],
+            }
+        }
 
-    async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         raise NotImplementedError
 
-    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+    async def alist(
+        self, config: RunnableConfig
+    ) -> AsyncIterator[tuple[RunnableConfig, Checkpoint]]:
+        raise NotImplementedError
+
+    async def aput(
+        self, config: RunnableConfig, checkpoint: Checkpoint
+    ) -> RunnableConfig:
         raise NotImplementedError