From 4271afeb4258baf01931d848df9b9394bd108fa2 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 26 Sep 2024 21:05:58 -0400 Subject: [PATCH] checkpoint: support serde for SecretStr --- .../checkpoint/langgraph/checkpoint/serde/jsonplus.py | 11 +++++++++++ libs/checkpoint/tests/test_jsonplus.py | 5 ++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 707442d34..d9796b0c3 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -233,6 +233,17 @@ def _msgpack_default(obj: Any) -> Union[str, msgpack.ExtType]: ), ), ) + elif hasattr(obj, "get_secret_value") and callable(obj.get_secret_value): + return msgpack.ExtType( + EXT_CONSTRUCTOR_SINGLE_ARG, + _msgpack_enc( + ( + obj.__class__.__module__, + obj.__class__.__name__, + obj.get_secret_value(), + ), + ), + ) elif hasattr(obj, "dict") and callable(obj.dict): # pydantic v1 return msgpack.ExtType( EXT_PYDANTIC_V1, diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index e52531414..1621a0c16 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -10,8 +10,9 @@ from ipaddress import IPv4Address import dataclasses_json -from pydantic import BaseModel +from pydantic import BaseModel, SecretStr from pydantic.v1 import BaseModel as BaseModelV1 +from pydantic.v1 import SecretStr as SecretStrV1 from zoneinfo import ZoneInfo from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer @@ -109,6 +110,8 @@ def test_serde_jsonplus() -> None: "my_pydantic_v1": MyPydanticV1( foo="foo", bar=1, inner=InnerPydanticV1(hello="hello") ), + "my_secret_str": SecretStr("meow"), + "my_secret_str_v1": SecretStrV1("meow"), "person": Person(name="foo"), "a_bool": True, "a_none": None,