mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 23:52:23 +02:00
Nuno Campos
74a40a895a
- 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
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
from collections import defaultdict
|
|
from typing import Any
|
|
|
|
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):
|
|
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"]
|
|
if saved := super().get(config):
|
|
assert self.storage_for_copies[thread_id][saved["ts"]] == saved
|
|
self.storage_for_copies[thread_id][checkpoint["ts"]] = copy_checkpoint(
|
|
checkpoint
|
|
)
|
|
# call super to write checkpoint
|
|
return super().put(config, checkpoint)
|