Files
langgraph/tests/memory_assert.py
T
Nuno Campos 52b2c755d7 Use uuid6 as the id for checkpoints
- this avoids conflicts if multiple processes creating checkpoints in same thread at same time
- uuid6 with a monotically increasing clock_seq is sortable by creation time at ms precision (plus clock_seq for ties, plus 48 bits of randomness for further ties)
2024-05-14 16:56:25 -07:00

49 lines
1.3 KiB
Python

from collections import defaultdict
from typing import Any, Optional
from langgraph.checkpoint.base import (
Checkpoint,
CheckpointMetadata,
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()
storage_for_copies: defaultdict[str, dict[str, Checkpoint]]
def __init__(
self,
*,
serde: Optional[SerializerProtocol] = None,
) -> None:
super().__init__(serde=serde)
self.storage_for_copies = defaultdict(dict)
def put(
self,
config: dict,
checkpoint: Checkpoint,
metadata: Optional[CheckpointMetadata] = None,
) -> 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["id"]] == saved
self.storage_for_copies[thread_id][checkpoint["id"]] = copy_checkpoint(
checkpoint
)
# call super to write checkpoint
return super().put(config, checkpoint, metadata)