diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py new file mode 100644 index 000000000..cda29dd27 --- /dev/null +++ b/langgraph/checkpoint/aiosqlite.py @@ -0,0 +1,86 @@ +import pickle +from contextlib import asynccontextmanager +from typing import Optional, final + +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() diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index 2d6aad224..75958708e 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/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 diff --git a/poetry.lock b/poetry.lock index 63903270c..e1154fe65 100644 --- a/poetry.lock +++ b/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 = "ca23b252b03db034e0da020db7769f4b1df7c49185ef2121d5db96d05b450af6" diff --git a/pyproject.toml b/pyproject.toml index 58808da0f..5a9c3574a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ jupyter = "^1.0.0" langchain = "^0.1.0" langchainhub = "^0.1.14" langchain-openai = "^0.0.2" +aiosqlite = "^0.19.0" [tool.ruff] select = [ "E", "F", "I" ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e829e2919..37a2f667f 100644 --- a/tests/test_pregel_async.py +++ b/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,57 @@ 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 + + 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))