Skip to content

Commit

Permalink
checkpoint: support serde for SecretStr
Browse files Browse the repository at this point in the history
  • Loading branch information
vbarda committed Sep 27, 2024
1 parent 29a0330 commit 4271afe
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 1 deletion.
11 changes: 11 additions & 0 deletions libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
5 changes: 4 additions & 1 deletion libs/checkpoint/tests/test_jsonplus.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 4271afe

Please sign in to comment.