Merge pull request #2735 from langchain-ai/vb/relax-strict-keys

checkpoint: set strict_map_key=False in serde
langchain-ai · Dec 12, 2024 · 2d8246e · 2d8246e
2 parents 083a14c + a031f82
commit 2d8246e
Showing 1 changed file with 18 additions and 6 deletions.
diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
@@ -438,28 +438,36 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]:
 def _msgpack_ext_hook(code: int, data: bytes) -> Any:
     if code == EXT_CONSTRUCTOR_SINGLE_ARG:
         try:
-            tup = msgpack.unpackb(data, ext_hook=_msgpack_ext_hook)
+            tup = msgpack.unpackb(
+                data, ext_hook=_msgpack_ext_hook, strict_map_key=False
+            )
             # module, name, arg
             return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
         except Exception:
             return
     elif code == EXT_CONSTRUCTOR_POS_ARGS:
         try:
-            tup = msgpack.unpackb(data, ext_hook=_msgpack_ext_hook)
+            tup = msgpack.unpackb(
+                data, ext_hook=_msgpack_ext_hook, strict_map_key=False
+            )
             # module, name, args
             return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
         except Exception:
             return
     elif code == EXT_CONSTRUCTOR_KW_ARGS:
         try:
-            tup = msgpack.unpackb(data, ext_hook=_msgpack_ext_hook)
+            tup = msgpack.unpackb(
+                data, ext_hook=_msgpack_ext_hook, strict_map_key=False
+            )
             # module, name, args
             return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
         except Exception:
             return
     elif code == EXT_METHOD_SINGLE_ARG:
         try:
-            tup = msgpack.unpackb(data, ext_hook=_msgpack_ext_hook)
+            tup = msgpack.unpackb(
+                data, ext_hook=_msgpack_ext_hook, strict_map_key=False
+            )
             # module, name, arg, method
             return getattr(getattr(importlib.import_module(tup[0]), tup[1]), tup[3])(
                 tup[2]
@@ -468,7 +476,9 @@ def _msgpack_ext_hook(code: int, data: bytes) -> Any:
             return
     elif code == EXT_PYDANTIC_V1:
         try:
-            tup = msgpack.unpackb(data, ext_hook=_msgpack_ext_hook)
+            tup = msgpack.unpackb(
+                data, ext_hook=_msgpack_ext_hook, strict_map_key=False
+            )
             # module, name, kwargs
             cls = getattr(importlib.import_module(tup[0]), tup[1])
             try:
@@ -479,7 +489,9 @@ def _msgpack_ext_hook(code: int, data: bytes) -> Any:
             return
     elif code == EXT_PYDANTIC_V2:
         try:
-            tup = msgpack.unpackb(data, ext_hook=_msgpack_ext_hook)
+            tup = msgpack.unpackb(
+                data, ext_hook=_msgpack_ext_hook, strict_map_key=False
+            )
             # module, name, kwargs, method
             cls = getattr(importlib.import_module(tup[0]), tup[1])
             try: