From f0fa5774662ec6b93c543ea8b343a4420fda06e8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jan 2024 16:23:21 -0800 Subject: [PATCH] Make ppickleable --- langgraph/checkpoint/__init__.py | 3 ++- langgraph/checkpoint/base.py | 6 +++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/langgraph/checkpoint/__init__.py b/langgraph/checkpoint/__init__.py index 7a52e0c6f..3d8a1ab76 100644 --- a/langgraph/checkpoint/__init__.py +++ b/langgraph/checkpoint/__init__.py @@ -1,7 +1,8 @@ -from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointAt +from langgraph.checkpoint.base import BaseCheckpointSaver, Checkpoint, CheckpointAt from langgraph.checkpoint.memory import MemorySaver __all__ = [ + "Checkpoint", "CheckpointAt", "BaseCheckpointSaver", "MemorySaver", diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index e40f3fcfc..699dc9dd8 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -19,13 +19,17 @@ class Checkpoint(TypedDict): versions_seen: defaultdict[str, defaultdict[str, int]] +def _seen_dict(): + return defaultdict(int) + + def empty_checkpoint() -> Checkpoint: return Checkpoint( v=1, ts=datetime.now(timezone.utc).isoformat(), channel_values={}, channel_versions=defaultdict(int), - versions_seen=defaultdict(lambda: defaultdict(int)), + versions_seen=defaultdict(_seen_dict), )