From 77f6ec5a2b7f165d136f66d7b01922a5ad2ce28e Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 13:03:30 -0400 Subject: [PATCH] update memory --- libs/langgraph/langgraph/checkpoint/memory.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index eec3fcdf7..f695d0f42 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -70,18 +70,20 @@ class MemorySaver(BaseCheckpointSaver): Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ thread_id = config["configurable"]["thread_id"] - if ts := config["configurable"].get("thread_ts"): - if saved := self.storage[thread_id].get(ts): - checkpoint, metadata = saved - writes = self.writes[(thread_id, ts)] - return CheckpointTuple( - config=config, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - pending_writes=[ - (id, c, self.serde.loads(v)) for id, c, v in writes - ], - ) + if thread_ts := config["configurable"].get("thread_ts"): + if checkpoints := self.storage[thread_id]: + ts = max(key for key in checkpoints.keys() if key <= thread_ts) + if saved := self.storage[thread_id].get(ts): + checkpoint, metadata = saved + writes = self.writes[(thread_id, ts)] + return CheckpointTuple( + config=config, + checkpoint=self.serde.loads(checkpoint), + metadata=self.serde.loads(metadata), + pending_writes=[ + (id, c, self.serde.loads(v)) for id, c, v in writes + ], + ) else: if checkpoints := self.storage[thread_id]: ts = max(checkpoints.keys())