From 083178d2a60209447bd882d3248214295be75ef6 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 1 Feb 2024 18:12:52 -0800
Subject: [PATCH] Add SqliteSaver

---
 langgraph/checkpoint/sqlite.py | 82 ++++++++++++++++++++++++++++++++++
 tests/test_pregel.py           | 51 +++++++++++++++++++++
 2 files changed, 133 insertions(+)
 create mode 100644 langgraph/checkpoint/sqlite.py

diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py
new file mode 100644
index 000000000..500606c29
--- /dev/null
+++ b/langgraph/checkpoint/sqlite.py
@@ -0,0 +1,82 @@
+import pickle
+import sqlite3
+from contextlib import contextmanager
+from typing import Optional
+
+from langchain_core.pydantic_v1 import Field
+from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables.utils import ConfigurableFieldSpec
+
+from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
+
+
+class SqliteSaver(BaseCheckpointSaver):
+    conn: sqlite3.Connection
+
+    is_setup: bool = Field(False, init=False, repr=False)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @classmethod
+    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 setup(self) -> None:
+        print("Setting up", self.is_setup)
+        if self.is_setup:
+            return
+
+        self.conn.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS checkpoints (
+                thread_id TEXT PRIMARY KEY,
+                checkpoint BLOB
+            );
+            """
+        )
+
+        self.is_setup = True
+
+    @contextmanager
+    def cursor(self, transaction: bool = True):
+        self.setup()
+        cur = self.conn.cursor()
+        try:
+            yield cur
+        finally:
+            if transaction:
+                self.conn.commit()
+            cur.close()
+
+    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
+        with self.cursor(transaction=False) as cur:
+            cur.execute(
+                "SELECT checkpoint FROM checkpoints WHERE thread_id = ?",
+                (config["configurable"]["thread_id"],),
+            )
+            if value := cur.fetchone():
+                return pickle.loads(value[0])
+
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+        with self.cursor() as cur:
+            cur.execute(
+                "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)",
+                (
+                    config["configurable"]["thread_id"],
+                    pickle.dumps(checkpoint),
+                ),
+            )
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 205c34d23..ebc04044a 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -16,6 +16,7 @@ from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
 from langgraph.checkpoint.memory import MemorySaver
+from langgraph.checkpoint.sqlite import SqliteSaver
 from langgraph.graph import END, Graph
 from langgraph.graph.message import MessageGraph
 from langgraph.graph.state import StateGraph
@@ -427,6 +428,56 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     assert checkpoint["channel_values"].get("total") == 5
 
 
+def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
+
+    def raise_if_above_10(input: int) -> int:
+        if input > 10:
+            raise ValueError("Input is too large")
+        return input
+
+    one = (
+        Channel.subscribe_to(["input"]).join(["total"])
+        | add_one
+        | Channel.write_to("output", "total")
+        | raise_if_above_10
+    )
+
+    memory = SqliteSaver.from_conn_string(":memory:")
+
+    app = Pregel(
+        nodes={"one": one},
+        channels={"total": BinaryOperatorAggregate(int, operator.add)},
+        checkpointer=memory,
+    )
+
+    # total starts out as 0, so output is 0+2=2
+    assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 2
+    # total is now 2, so output is 2+3=5
+    assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 7
+    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+    with pytest.raises(ValueError):
+        app.invoke(4, {"configurable": {"thread_id": "1"}})
+    # checkpoint is not updated
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 7
+    # on a new thread, total starts out as 0, so output is 0+5=5
+    assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
+    checkpoint = memory.get({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 7
+    checkpoint = memory.get({"configurable": {"thread_id": "2"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 5
+
+
 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))