From 74a40a895a3363b3396f50ae3b19a3774984ce99 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 12 Apr 2024 11:49:29 -0700 Subject: [PATCH 1/2] Make ser/de configurable in checkpointer classes - Remove pydantic usage from base checkpointer class - Make serialization configurable for all existing checkpointer classes, you can eg use dill or json instead of pickle --- langgraph/checkpoint/__init__.py | 10 +++- langgraph/checkpoint/aiosqlite.py | 33 ++++++++---- langgraph/checkpoint/base.py | 54 ++++++++++++------- langgraph/checkpoint/memory.py | 87 ++++++++++++++++++++++--------- langgraph/checkpoint/sqlite.py | 48 +++++++++-------- tests/memory_assert.py | 30 ++++++++--- 6 files changed, 177 insertions(+), 85 deletions(-) diff --git a/langgraph/checkpoint/__init__.py b/langgraph/checkpoint/__init__.py index 3d8a1ab76..5a0f4bd0c 100644 --- a/langgraph/checkpoint/__init__.py +++ b/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", ] diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index ada5c10d1..b508cefb9 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/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() diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index e63a00201..6e3b2fc46 100644 --- a/langgraph/checkpoint/base.py +++ b/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 diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 059c58573..c3348ed29 100644 --- a/langgraph/checkpoint/memory.py +++ b/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 + ) diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index 3a292fb5a..829874646 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/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") diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 50c185a12..2fb3a0127 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -1,18 +1,36 @@ from collections import defaultdict +from typing import Any -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: SerializerProtocol | None = None, at: CheckpointAt | None = 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"] From f5a7c01e886b1b78279e681a4a551419ed8a2708 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 12 Apr 2024 11:55:20 -0700 Subject: [PATCH 2/2] py39 --- tests/memory_assert.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 2fb3a0127..6bceb865e 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -1,5 +1,5 @@ from collections import defaultdict -from typing import Any +from typing import Any, Optional from langgraph.checkpoint.base import ( Checkpoint, @@ -26,7 +26,10 @@ class MemorySaverAssertImmutable(MemorySaver): storage_for_copies: defaultdict[str, dict[str, Checkpoint]] def __init__( - self, *, serde: SerializerProtocol | None = None, at: CheckpointAt | None = None + self, + *, + serde: Optional[SerializerProtocol] = None, + at: Optional[CheckpointAt] = None, ) -> None: super().__init__(serde=serde, at=at) self.storage_for_copies = defaultdict(dict)