Merge pull request #305 from langchain-ai/nc/12apr/configurable-serde

Make ser/de configurable in checkpointer classes

langgraph/checkpoint/__init__.py

@@ -1,9 +1,15 @@
-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointAt
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointAt,
+    SerializerProtocol,
+)
 from langgraph.checkpoint.memory import MemorySaver
 
 __all__ = [
+    "BaseCheckpointSaver",
     "Checkpoint",
     "CheckpointAt",
-    "BaseCheckpointSaver",
     "MemorySaver",
+    "SerializerProtocol",
 ]

langgraph/checkpoint/aiosqlite.py

@@ -4,20 +4,35 @@ from types import TracebackType
 from typing import AsyncIterator, Optional
 
 import aiosqlite
-from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointAt,
+    CheckpointTuple,
+    SerializerProtocol,
+)
 
 
 class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
+    serde = pickle
+
     conn: aiosqlite.Connection
 
-    is_setup: bool = Field(False, init=False, repr=False)
+    is_setup: bool
 
-    class Config:
-        arbitrary_types_allowed = True
+    def __init__(
+        self,
+        conn: aiosqlite.Connection,
+        *,
+        serde: Optional[SerializerProtocol] = None,
+        at: Optional[CheckpointAt] = None,
+    ):
+        super().__init__(serde=serde, at=at)
+        self.conn = conn
+        self.is_setup = False
 
     @classmethod
     def from_conn_string(cls, conn_string: str) -> "AsyncSqliteSaver":
@@ -67,7 +82,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
                 if value := await cursor.fetchone():
                     return CheckpointTuple(
                         config,
-                        pickle.loads(value[0]),
+                        self.serde.loads(value[0]),
                         {
                             "configurable": {
                                 "thread_id": config["configurable"]["thread_id"],
@@ -90,7 +105,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
                                 "thread_ts": value[1],
                             }
                         },
-                        pickle.loads(value[3]),
+                        self.serde.loads(value[3]),
                         {
                             "configurable": {
                                 "thread_id": value[0],
@@ -110,7 +125,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
             async for thread_id, thread_ts, parent_ts, value in cursor:
                 yield CheckpointTuple(
                     {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
-                    pickle.loads(value),
+                    self.serde.loads(value),
                     {"configurable": {"thread_id": thread_id, "thread_ts": parent_ts}}
                     if parent_ts
                     else None,
@@ -126,7 +141,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
                 config["configurable"]["thread_id"],
                 checkpoint["ts"],
                 config["configurable"].get("thread_ts"),
-                pickle.dumps(checkpoint),
+                self.serde.dumps(checkpoint),
             ),
         ):
             await self.conn.commit()

langgraph/checkpoint/base.py

@@ -1,13 +1,18 @@
-import asyncio
 from abc import ABC
 from collections import defaultdict
 from copy import deepcopy
 from datetime import datetime, timezone
-from typing import Any, AsyncIterator, Iterator, NamedTuple, Optional, TypedDict
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterator,
+    NamedTuple,
+    Optional,
+    Protocol,
+    TypedDict,
+)
 
-from langchain_core.load.serializable import Serializable
-from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.utils import ConfigurableFieldSpec
+from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
 
 from langgraph.utils import StrEnum
 
@@ -93,8 +98,27 @@ CheckpointThreadTs = ConfigurableFieldSpec(
 )
 
 
-class BaseCheckpointSaver(Serializable, ABC):
-    at: CheckpointAt = CheckpointAt.END_OF_RUN
+class SerializerProtocol(Protocol):
+    def dumps(self, obj: Any) -> bytes:
+        ...
+
+    def loads(self, data: bytes) -> Any:
+        ...
+
+
+class BaseCheckpointSaver(ABC):
+    at: CheckpointAt = CheckpointAt.END_OF_STEP
+
+    serde: SerializerProtocol
+
+    def __init__(
+        self,
+        *,
+        serde: Optional[SerializerProtocol] = None,
+        at: Optional[CheckpointAt] = None,
+    ) -> None:
+        self.serde = serde or self.serde
+        self.at = at or self.at
 
     @property
     def config_specs(self) -> list[ConfigurableFieldSpec]:
@@ -118,22 +142,12 @@ class BaseCheckpointSaver(Serializable, ABC):
             return value.checkpoint
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.get_tuple, config
-        )
+        raise NotImplementedError
 
     async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]:
-        loop = asyncio.get_running_loop()
-        iter = loop.run_in_executor(None, self.list, config)
-        while True:
-            try:
-                yield await loop.run_in_executor(None, next, iter)
-            except StopIteration:
-                return
+        raise NotImplementedError
 
     async def aput(
         self, config: RunnableConfig, checkpoint: Checkpoint
     ) -> RunnableConfig:
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.put, config, checkpoint
-        )
+        raise NotImplementedError

langgraph/checkpoint/memory.py

@@ -1,43 +1,59 @@
+import asyncio
+import pickle
 from collections import defaultdict
-from typing import Optional
+from typing import AsyncIterator, Iterator, Optional
 
-from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
 
-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointAt,
+    CheckpointTuple,
+    SerializerProtocol,
+)
 
 
 class MemorySaver(BaseCheckpointSaver):
-    storage: defaultdict[str, dict[str, Checkpoint]] = Field(
-        default_factory=lambda: defaultdict(dict)
-    )
+    serde = pickle
 
-    def get(self, config: RunnableConfig) -> Optional[Checkpoint]:
-        if value := self.get_tuple(config):
-            return value.checkpoint
+    storage: defaultdict[str, dict[str, Checkpoint]]
+
+    def __init__(
+        self,
+        *,
+        serde: Optional[SerializerProtocol] = None,
+        at: Optional[CheckpointAt] = None,
+    ) -> None:
+        super().__init__(serde=serde, at=at)
+        self.storage = defaultdict(dict)
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
-        if config["configurable"].get("thread_ts"):
-            if checkpoint := self.storage[config["configurable"]["thread_id"]].get(
-                config["configurable"]["thread_ts"]
-            ):
-                return CheckpointTuple(config=config, checkpoint=checkpoint)
-        else:
-            if checkpoints := self.storage[config["configurable"]["thread_id"]]:
-                thread_ts = max(checkpoints.keys())
+        thread_id = config["configurable"]["thread_id"]
+        if ts := config["configurable"].get("thread_ts"):
+            if checkpoint := self.storage[thread_id].get(ts):
                 return CheckpointTuple(
-                    config={
-                        "configurable": {
-                            "thread_id": config["configurable"]["thread_id"],
-                            "thread_ts": thread_ts,
-                        }
-                    },
-                    checkpoint=checkpoints[thread_ts],
+                    config=config, checkpoint=self.serde.loads(checkpoint)
                 )
+        else:
+            if checkpoints := self.storage[thread_id]:
+                ts = max(checkpoints.keys())
+                return CheckpointTuple(
+                    config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
+                    checkpoint=self.serde.loads(checkpoints[ts]),
+                )
+
+    def list(self, config: RunnableConfig) -> Iterator[CheckpointTuple]:
+        thread_id = config["configurable"]["thread_id"]
+        for ts, checkpoint in self.storage[thread_id].items():
+            yield CheckpointTuple(
+                config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
+                checkpoint=self.serde.loads(checkpoint),
+            )
 
     def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> RunnableConfig:
         self.storage[config["configurable"]["thread_id"]].update(
-            {checkpoint["ts"]: checkpoint}
+            {checkpoint["ts"]: self.serde.dumps(checkpoint)}
         )
         return {
             "configurable": {
@@ -45,3 +61,24 @@ class MemorySaver(BaseCheckpointSaver):
                 "thread_ts": checkpoint["ts"],
             }
         }
+
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.get_tuple, config
+        )
+
+    async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]:
+        loop = asyncio.get_running_loop()
+        iter = loop.run_in_executor(None, self.list, config)
+        while True:
+            try:
+                yield await loop.run_in_executor(None, next, iter)
+            except StopIteration:
+                return
+
+    async def aput(
+        self, config: RunnableConfig, checkpoint: Checkpoint
+    ) -> RunnableConfig:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.put, config, checkpoint
+        )

langgraph/checkpoint/sqlite.py

@@ -2,22 +2,37 @@ import pickle
 import sqlite3
 from contextlib import AbstractContextManager, contextmanager
 from types import TracebackType
-from typing import AsyncIterator, Iterator, Optional
+from typing import Iterator, Optional
 
-from langchain_core.pydantic_v1 import Field
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
-from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointTuple
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointAt,
+    CheckpointTuple,
+    SerializerProtocol,
+)
 
 
 class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
+    serde = pickle
+
     conn: sqlite3.Connection
 
-    is_setup: bool = Field(False, init=False, repr=False)
+    is_setup: bool
 
-    class Config:
-        arbitrary_types_allowed = True
+    def __init__(
+        self,
+        conn: sqlite3.Connection,
+        *,
+        serde: Optional[SerializerProtocol] = None,
+        at: Optional[CheckpointAt] = None,
+    ) -> None:
+        super().__init__(serde=serde, at=at)
+        self.conn = conn
+        self.is_setup = False
 
     @classmethod
     def from_conn_string(cls, conn_string: str) -> "SqliteSaver":
@@ -76,7 +91,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                 if value := cur.fetchone():
                     return CheckpointTuple(
                         config,
-                        pickle.loads(value[0]),
+                        self.serde.loads(value[0]),
                         {
                             "configurable": {
                                 "thread_id": config["configurable"]["thread_id"],
@@ -99,7 +114,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                                 "thread_ts": value[1],
                             }
                         },
-                        pickle.loads(value[3]),
+                        self.serde.loads(value[3]),
                         {
                             "configurable": {
                                 "thread_id": value[0],
@@ -119,7 +134,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
             for thread_id, thread_ts, parent_ts, value in cur:
                 yield CheckpointTuple(
                     {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
-                    pickle.loads(value),
+                    self.serde.loads(value),
                     {
                         "configurable": {
                             "thread_id": thread_id,
@@ -138,7 +153,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                     config["configurable"]["thread_id"],
                     checkpoint["ts"],
                     config["configurable"].get("thread_ts"),
-                    pickle.dumps(checkpoint),
+                    self.serde.dumps(checkpoint),
                 ),
             )
         return {
@@ -147,16 +162,3 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
                 "thread_ts": checkpoint["ts"],
             }
         }
-
-    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
-        raise NotImplementedError("Use AsyncSqliteSaver instead")
-
-    async def alist(
-        self, config: RunnableConfig
-    ) -> AsyncIterator[tuple[RunnableConfig, Checkpoint]]:
-        raise NotImplementedError("Use AsyncSqliteSaver instead")
-
-    async def aput(
-        self, config: RunnableConfig, checkpoint: Checkpoint
-    ) -> RunnableConfig:
-        raise NotImplementedError("Use AsyncSqliteSaver instead")

tests/memory_assert.py

@@ -1,18 +1,39 @@
 from collections import defaultdict
+from typing import Any, Optional
 
-from langchain_core.pydantic_v1 import Field
-
-from langgraph.checkpoint.base import Checkpoint, CheckpointAt, copy_checkpoint
+from langgraph.checkpoint.base import (
+    Checkpoint,
+    CheckpointAt,
+    SerializerProtocol,
+    copy_checkpoint,
+)
 from langgraph.checkpoint.memory import MemorySaver
 
 
+class NoopSerializer(SerializerProtocol):
+    def loads(self, data: bytes) -> Any:
+        return data
+
+    def dumps(self, obj: Any) -> bytes:
+        return obj
+
+
 class MemorySaverAssertImmutable(MemorySaver):
-    storage_for_copies: defaultdict[str, dict[str, Checkpoint]] = Field(
-        default_factory=lambda: defaultdict(dict)
-    )
+    serde = NoopSerializer()
 
     at = CheckpointAt.END_OF_STEP
 
+    storage_for_copies: defaultdict[str, dict[str, Checkpoint]]
+
+    def __init__(
+        self,
+        *,
+        serde: Optional[SerializerProtocol] = None,
+        at: Optional[CheckpointAt] = None,
+    ) -> None:
+        super().__init__(serde=serde, at=at)
+        self.storage_for_copies = defaultdict(dict)
+
     def put(self, config: dict, checkpoint: Checkpoint) -> None:
         # assert checkpoint hasn't been modified since last written
         thread_id = config["configurable"]["thread_id"]