From d98e90a13ecfd2ff802d233fcc3969a07ca320bd Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 13 Jun 2024 18:14:30 -0700
Subject: [PATCH] Support float and str versions

---
 langgraph/checkpoint/base.py | 12 ++++++++----
 langgraph/pregel/__init__.py |  6 +++---
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py
index 3ef13b23f..98f464731 100644
--- a/langgraph/checkpoint/base.py
+++ b/langgraph/checkpoint/base.py
@@ -11,6 +11,8 @@ from typing import (
     NamedTuple,
     Optional,
     TypedDict,
+    TypeVar,
+    Union,
 )
 
 from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
@@ -21,6 +23,8 @@ from langgraph.constants import Send
 from langgraph.serde.base import SerializerProtocol
 from langgraph.serde.jsonplus import JsonPlusSerializer
 
+V = TypeVar("V", int, float, str)
+
 
 # Marked as total=False to allow for future expansion.
 class CheckpointMetadata(TypedDict, total=False):
@@ -63,13 +67,13 @@ class Checkpoint(TypedDict):
     
     Mapping from channel name to channel snapshot value.
     """
-    channel_versions: defaultdict[str, int]
+    channel_versions: defaultdict[str, Union[str, int, float]]
     """The versions of the channels at the time of the checkpoint.
     
     The keys are channel names and the values are the logical time step
     at which the channel was last updated.
     """
-    versions_seen: defaultdict[str, defaultdict[str, int]]
+    versions_seen: defaultdict[str, defaultdict[str, Union[str, int, float]]]
     """Map from node ID to map from channel name to version seen.
     
     This keeps track of the versions of the channels that each node has seen.
@@ -203,5 +207,5 @@ class BaseCheckpointSaver(ABC):
     ) -> RunnableConfig:
         raise NotImplementedError
 
-    def get_next_version(self, current: int, channel: BaseChannel) -> int:
-        return current + 1
+    def get_next_version(self, current: Optional[V], channel: BaseChannel) -> V:
+        return current + 1 if current is not None else 1
diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py
index f8fd0facc..f8e86206f 100644
--- a/langgraph/pregel/__init__.py
+++ b/langgraph/pregel/__init__.py
@@ -1593,8 +1593,8 @@ def _local_write(
     commit(writes)
 
 
-def _increment(current: int, channel: BaseChannel) -> int:
-    return current + 1
+def _increment(current: Optional[int], channel: BaseChannel) -> int:
+    return current + 1 if current is not None else 1
 
 
 def _apply_writes(
@@ -1631,7 +1631,7 @@ def _apply_writes(
                     f"Invalid update for channel {chan} with values {vals}"
                 ) from e
             checkpoint["channel_versions"][chan] = get_next_version(
-                max_version, channels[chan]
+                max_version or None, channels[chan]
             )
             updated_channels.add(chan)
     # Channels that weren't updated in this step are notified of a new step