WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-22 01:25:06 +02:00
Code Issues Packages Projects Releases Wiki Activity

checkpoint: switch thread_ts -> checkpoint_id, add checkpoint_ns, change serializer protocol (#1185)

Browse Source 
---------

Co-authored-by: Nuno Campos <nuno@langchain.dev>
This commit is contained in:
Vadym Barda
2024-08-02 01:08:19 +00:00
committed by GitHub
co-authored by Nuno Campos
parent 862afa27de
commit 4b2187c9a3
22 changed files with 1491 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/tests/memory_assert.py
+22 -18
View File
@@ -15,17 +15,17 @@ from langgraph.checkpoint.memory import MemorySaver
class NoopSerializer(SerializerProtocol):
def loads(self, data: bytes) -> Any:
return data
def loads_typed(self, data: tuple[str, bytes]) -> Any:
return data[1]
def dumps(self, obj: Any) -> bytes:
return obj
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
return "type", obj
class MemorySaverAssertImmutable(MemorySaver):
serde = NoopSerializer()
storage_for_copies: defaultdict[str, dict[str, Checkpoint]]
storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]]
def __init__(
self,
@@ -34,7 +34,7 @@ class MemorySaverAssertImmutable(MemorySaver):
put_sleep: Optional[float] = None,
) -> None:
super().__init__(serde=serde)
self.storage_for_copies = defaultdict(dict)
self.storage_for_copies = defaultdict(lambda: defaultdict(dict))
self.put_sleep = put_sleep
def put(
@@ -49,14 +49,17 @@ class MemorySaverAssertImmutable(MemorySaver):
time.sleep(self.put_sleep)
# assert checkpoint hasn't been modified since last written
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
if saved := super().get(config):
assert (
self.serde.loads(self.storage_for_copies[thread_id][saved["id"]])
self.serde.loads_typed(
self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]]
)
== saved
)
self.storage_for_copies[thread_id][checkpoint["id"]] = self.serde.dumps(
copy_checkpoint(checkpoint)
)
self.storage_for_copies[thread_id][checkpoint_ns][
checkpoint["id"]
] = self.serde.dumps_typed(copy_checkpoint(checkpoint))
# call super to write checkpoint
return super().put(config, checkpoint, metadata)
@@ -91,23 +94,24 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver):
"""
configurable = config["configurable"].copy()
# remove thread_ts to make testing simpler
thread_ts = configurable.pop("thread_ts", None)
self.storage[config["configurable"]["thread_id"]].update(
# remove checkpoint_id to make testing simpler
checkpoint_id = configurable.pop("checkpoint_id", None)
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"]["checkpoint_ns"]
self.storage[thread_id][checkpoint_ns].update(
{
checkpoint["id"]: (
self.serde.dumps(checkpoint),
self.serde.dumps_typed(checkpoint),
# merge configurable fields and metadata
self.serde.dumps({**configurable, **metadata}),
thread_ts,
self.serde.dumps_typed({**configurable, **metadata}),
checkpoint_id,
)
}
)
return {
"configurable": {
"thread_id": config["configurable"]["thread_id"],
"thread_ts": checkpoint["id"],
"checkpoint_id": checkpoint["id"],
}
}