Merge pull request #91 from langchain-ai/nc/7feb/aiosqlite

Add AsyncSqliteSaver

Makefile

@@ -18,7 +18,7 @@ test:
 	poetry run pytest
 
 test_watch:
-	poetry run ptw --snapshot-update --now . -- -vv -x tests
+	poetry run ptw --snapshot-update --now . -- -vv -x --ff tests
 
 ######################
 # LINTING AND FORMATTING

langgraph/checkpoint/aiosqlite.py

@@ -0,0 +1,85 @@
+import pickle
+from typing import Optional
+
+import aiosqlite
+from langchain_core.pydantic_v1 import Field
+from langchain_core.runnables import RunnableConfig
+from langchain_core.runnables.utils import ConfigurableFieldSpec
+
+from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint
+
+
+class AsyncSqliteSaver(BaseCheckpointSaver):
+    conn: aiosqlite.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) -> "AsyncSqliteSaver":
+        return AsyncSqliteSaver(conn=aiosqlite.connect(conn_string))
+
+    @property
+    def config_specs(self) -> list[ConfigurableFieldSpec]:
+        return [
+            ConfigurableFieldSpec(
+                id="thread_id",
+                annotation=str,
+                name="Thread ID",
+                description=None,
+                default="",
+                is_shared=True,
+            ),
+        ]
+
+    async def setup(self) -> None:
+        print("hello")
+        if self.is_setup:
+            return
+
+        try:
+            await self.conn
+            await self.conn.executescript(
+                """
+                CREATE TABLE IF NOT EXISTS checkpoints (
+                    thread_id TEXT PRIMARY KEY,
+                    checkpoint BLOB
+                );
+                """
+            )
+            await self.conn.commit()
+
+            print("good bye")
+
+            self.is_setup = True
+        except BaseException as e:
+            print(e)
+            raise e
+
+    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
+        raise NotImplementedError
+
+    def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+        raise NotImplementedError
+
+    async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
+        await self.setup()
+        async with self.conn.execute(
+            "SELECT checkpoint FROM checkpoints WHERE thread_id = ?",
+            (config["configurable"]["thread_id"],),
+        ) as cursor:
+            if value := await cursor.fetchone():
+                return pickle.loads(value[0])
+
+    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+        await self.setup()
+        await self.conn.execute(
+            "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint) VALUES (?, ?)",
+            (
+                config["configurable"]["thread_id"],
+                pickle.dumps(checkpoint),
+            ),
+        )
+        await self.conn.commit()

langgraph/checkpoint/sqlite.py

@@ -79,3 +79,9 @@ class SqliteSaver(BaseCheckpointSaver):
                     pickle.dumps(checkpoint),
                 ),
             )
+
+    async def aget(self, config: RunnableConfig) -> Optional[Checkpoint]:
+        raise NotImplementedError
+
+    async def aput(self, config: RunnableConfig, checkpoint: Checkpoint) -> None:
+        raise NotImplementedError

langgraph/pregel/__init__.py

@@ -392,11 +392,10 @@ class Pregel(
         finally:
             # cancel any pending tasks when generator is interrupted
             try:
-                futures
+                for task in futures:
+                    task.cancel()
             except NameError:
-                return
-            for task in futures:
-                task.cancel()
+                pass
 
     async def _atransform(
         self,
@@ -562,11 +561,10 @@ class Pregel(
         finally:
             # cancel any pending tasks when generator is interrupted
             try:
-                futures
+                for task in futures:
+                    task.cancel()
             except NameError:
-                return
-            for task in futures:
-                task.cancel()
+                pass
 
     def invoke(
         self,

poetry.lock

@@ -110,6 +110,21 @@ files = [
 [package.dependencies]
 frozenlist = ">=1.1.0"
 
+[[package]]
+name = "aiosqlite"
+version = "0.19.0"
+description = "asyncio bridge to the standard sqlite3 module"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "aiosqlite-0.19.0-py3-none-any.whl", hash = "sha256:edba222e03453e094a3ce605db1b970c4b3376264e56f32e2a4959f948d66a96"},
+    {file = "aiosqlite-0.19.0.tar.gz", hash = "sha256:95ee77b91c8d2808bd08a59fbebf66270e9090c3d92ffbf260dc0db0b979577d"},
+]
+
+[package.extras]
+dev = ["aiounittest (==1.4.1)", "attribution (==1.6.2)", "black (==23.3.0)", "coverage[toml] (==7.2.3)", "flake8 (==5.0.4)", "flake8-bugbear (==23.3.12)", "flit (==3.7.1)", "mypy (==1.2.0)", "ufmt (==2.1.0)", "usort (==1.0.6)"]
+docs = ["sphinx (==6.1.3)", "sphinx-mdinclude (==0.5.3)"]
+
 [[package]]
 name = "annotated-types"
 version = "0.6.0"
@@ -3714,4 +3729,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "faf7cebfb8e64c2edefb69bacdbdd10d8399ea8b3ba9f46c4383e72bd3cd925a"
+content-hash = "f489f2e9159e8db255a43027617c41a389afb5fe84ff94bf8521c0f98de8cd38"

pyproject.toml

@@ -25,6 +25,7 @@ syrupy = "^4.0.2"
 httpx = "^0.26.0"
 pytest-watcher = "^0.3.4"
 langchain = "^0.1.0"
+aiosqlite = "^0.19.0"
 
 [tool.poetry.group.lint.dependencies]
 ruff = "^0.1.4"

tests/test_pregel_async.py

@@ -22,6 +22,7 @@ from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
+from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph
@@ -476,6 +477,59 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     assert checkpoint["channel_values"].get("total") == 5
 
 
+async 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 = AsyncSqliteSaver.from_conn_string(":memory:")
+
+    app = Pregel(
+        nodes={"one": one},
+        channels={"total": BinaryOperatorAggregate(int, operator.add)},
+        checkpointer=memory,
+        debug=True,
+    )
+
+    # total starts out as 0, so output is 0+2=2
+    assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 2
+    # total is now 2, so output is 2+3=5
+    assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 7
+    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+    with pytest.raises(ValueError):
+        await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
+    # checkpoint is not updated
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 7
+    # on a new thread, total starts out as 0, so output is 0+5=5
+    assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
+    checkpoint = await memory.aget({"configurable": {"thread_id": "1"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 7
+    checkpoint = await memory.aget({"configurable": {"thread_id": "2"}})
+    assert checkpoint is not None
+    assert checkpoint["channel_values"].get("total") == 5
+
+    await memory.conn.close()
+
+
 async 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))