From 728b2830639802e54018b5a8479725ae96981047 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 14 Aug 2024 10:59:14 -0700
Subject: [PATCH] checkpoint: Ignore serialized values when constructor no
 longer available

---
 .../langgraph/checkpoint/serde/jsonplus.py    | 27 ++++++++++---------
 libs/checkpoint/tests/test_jsonplus.py        | 18 +++++++++++++
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
index adbe98d33..52e9d6da3 100644
--- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
@@ -121,18 +121,21 @@ class JsonPlusSerializer(SerializerProtocol):
             and value.get("type", None) == "constructor"
             and value.get("id", None) is not None
         ):
-            # Get module and class name
-            [*module, name] = value["id"]
-            # Import module
-            mod = importlib.import_module(".".join(module))
-            # Import class
-            cls = getattr(mod, name)
-            # Instantiate class
-            if value["method"] is not None:
-                method = getattr(cls, value["method"])
-                return method(*value["args"], **value["kwargs"])
-            else:
-                return cls(*value["args"], **value["kwargs"])
+            try:
+                # Get module and class name
+                [*module, name] = value["id"]
+                # Import module
+                mod = importlib.import_module(".".join(module))
+                # Import class
+                cls = getattr(mod, name)
+                # Instantiate class
+                if value["method"] is not None:
+                    method = getattr(cls, value["method"])
+                    return method(*value["args"], **value["kwargs"])
+                else:
+                    return cls(*value["args"], **value["kwargs"])
+            except (ImportError, AttributeError):
+                return None
 
         return LC_REVIVER(value)
 
diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py
index 2ca5a59c1..d5a53d248 100644
--- a/libs/checkpoint/tests/test_jsonplus.py
+++ b/libs/checkpoint/tests/test_jsonplus.py
@@ -149,3 +149,21 @@ def test_serde_jsonplus_bytearray() -> None:
 
     assert dumped == ("bytearray", some_bytearray)
     assert serde.loads_typed(dumped) == some_bytearray
+
+
+def test_loads_cannot_find() -> None:
+    serde = JsonPlusSerializer()
+
+    dumped = (
+        "json",
+        b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydanticccc"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
+    )
+
+    assert serde.loads_typed(dumped) is None, "Should return None if cannot find class"
+
+    dumped = (
+        "json",
+        b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonpluss", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
+    )
+
+    assert serde.loads_typed(dumped) is None, "Should return None if cannot find module"