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
This commit is contained in:
Nuno Campos
2024-04-12 11:49:29 -07:00
parent aa07d30e1f
commit 74a40a895a
6 changed files with 177 additions and 85 deletions
+24 -6
View File
@@ -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"]