mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-02 14:28:46 +02:00
Add optional encryption of checkpointer payloads
- no dependency on any particular encryption lib (there is no py stdlib encryption lib) - works with any modern checkpointer, ie. those which use dumps_typed and loads_typed methods to serialize data - backwards compatible with unencrypted data in same storage (will just be read unencrypted) - providing easy constructor to use AES encryption through pycriptodome library, one single line of code to add it in - other encryption libraries or algorithms (even assymetric ones) can be used by implementing the two-method CipherProtocol interface - cipher name (eg. aes) is stored with encrypted payload for forwards compatibility
This commit is contained in:
Nuno Campos
@@ -45,3 +45,22 @@ def maybe_add_typed_methods(serde: SerializerProtocol) -> SerializerProtocol:
|
||||
return SerializerCompat(serde)
|
||||
|
||||
return serde
|
||||
|
||||
|
||||
class CipherProtocol(Protocol):
|
||||
"""Protocol for encryption and decryption of data.
|
||||
- `encrypt`: Encrypt plaintext.
|
||||
- `decrypt`: Decrypt ciphertext.
|
||||
- `name`: Name of the cipher.
|
||||
"""
|
||||
|
||||
name: str
|
||||
"""Name of the cipher."""
|
||||
|
||||
def encrypt(self, plaintext: bytes) -> bytes:
|
||||
"""Encrypt plaintext."""
|
||||
...
|
||||
|
||||
def decrypt(self, ciphertext: bytes) -> bytes:
|
||||
"""Decrypt ciphertext."""
|
||||
...
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from langgraph.checkpoint.serde.base import CipherProtocol
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
|
||||
class EncryptedSerializer(JsonPlusSerializer):
|
||||
"""Serializer that encrypts and decrypts data using an encryption protocol."""
|
||||
|
||||
def __init__(self, cipher: CipherProtocol) -> None:
|
||||
super().__init__()
|
||||
self.cipher = cipher
|
||||
|
||||
def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
|
||||
"""Serialize an object to a tuple (type, bytes) and encrypt the bytes."""
|
||||
# serialize data
|
||||
typ, data = super().dumps_typed(obj)
|
||||
# encrypt data
|
||||
encrypted_data = self.cipher.encrypt(data, self.key)
|
||||
# add cipher name to type
|
||||
return f"{typ}+{self.cipher.name}", encrypted_data
|
||||
|
||||
def loads_typed(self, data: tuple[str, bytes]) -> Any:
|
||||
enc_typ, enc_data = data
|
||||
# unencrypted data
|
||||
if "+" not in enc_typ:
|
||||
return super().loads_typed(data)
|
||||
# verify cipher name
|
||||
typ, cipher_name = enc_typ.split("+", 1)
|
||||
if cipher_name != self.cipher.name:
|
||||
raise ValueError(
|
||||
f"Cipher mismatch: expected {self.cipher.name}, got {cipher_name}"
|
||||
)
|
||||
# decrypt data
|
||||
decrypted_data = self.cipher.decrypt(enc_data, self.key)
|
||||
# deserialize data
|
||||
return super().loads_typed((typ, decrypted_data))
|
||||
|
||||
@classmethod
|
||||
def from_pycryptodome_aes(cls, **kwargs: Any) -> "EncryptedSerializer":
|
||||
"""Create an EncryptedSerializer using AES encryption."""
|
||||
try:
|
||||
from Crypto.Cipher import AES # type: ignore
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Pycryptodome is not installed. Please install it with `pip install pycryptodome`."
|
||||
) from None
|
||||
|
||||
# check if AES key is provided
|
||||
if "key" in kwargs:
|
||||
key = kwargs.pop("key")
|
||||
else:
|
||||
key = os.getenvb("LANGGRAPH_AES_KEY")
|
||||
if key is None:
|
||||
raise ValueError("LANGGRAPH_AES_KEY environment variable is not set.")
|
||||
if len(key) not in (16, 24, 32):
|
||||
raise ValueError("LANGGRAPH_AES_KEY must be 16, 24, or 32 bytes long.")
|
||||
|
||||
# set default mode to EAX if not provided
|
||||
if kwargs.get("mode") is None:
|
||||
kwargs["mode"] = AES.MODE_EAX
|
||||
|
||||
class PycryptodomeAesCipher(CipherProtocol):
|
||||
name: str = "aes"
|
||||
|
||||
def encrypt(self, plaintext: bytes) -> bytes:
|
||||
cipher = AES.new(key, **kwargs)
|
||||
ciphertext, tag = cipher.encrypt_and_digest(plaintext)
|
||||
return cipher.nonce + tag + ciphertext
|
||||
|
||||
def decrypt(self, ciphertext: bytes) -> bytes:
|
||||
nonce = ciphertext[:16]
|
||||
tag = ciphertext[16:32]
|
||||
actual_ciphertext = ciphertext[32:]
|
||||
|
||||
cipher = AES.new(key, **kwargs, nonce=nonce)
|
||||
return cipher.decrypt_and_verify(actual_ciphertext, tag)
|
||||
|
||||
return cls(PycryptodomeAesCipher())
|
||||
Generated
+44
-2
@@ -1,4 +1,4 @@
|
||||
# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
|
||||
# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
|
||||
|
||||
[[package]]
|
||||
name = "aiosqlite"
|
||||
@@ -2238,6 +2238,48 @@ files = [
|
||||
{file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pycryptodome"
|
||||
version = "3.21.0"
|
||||
description = "Cryptographic library for Python"
|
||||
optional = false
|
||||
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
|
||||
groups = ["dev"]
|
||||
files = [
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:dad9bf36eda068e89059d1f07408e397856be9511d7113ea4b586642a429a4fd"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:a1752eca64c60852f38bb29e2c86fca30d7672c024128ef5d70cc15868fa10f4"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:3ba4cc304eac4d4d458f508d4955a88ba25026890e8abff9b60404f76a62c55e"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7cb087b8612c8a1a14cf37dd754685be9a8d9869bed2ffaaceb04850a8aeef7e"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-musllinux_1_1_aarch64.whl", hash = "sha256:26412b21df30b2861424a6c6d5b1d8ca8107612a4cfa4d0183e71c5d200fb34a"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-win32.whl", hash = "sha256:cc2269ab4bce40b027b49663d61d816903a4bd90ad88cb99ed561aadb3888dd3"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27m-win_amd64.whl", hash = "sha256:0fa0a05a6a697ccbf2a12cec3d6d2650b50881899b845fac6e87416f8cb7e87d"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:6cce52e196a5f1d6797ff7946cdff2038d3b5f0aba4a43cb6bf46b575fd1b5bb"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:a915597ffccabe902e7090e199a7bf7a381c5506a747d5e9d27ba55197a2c568"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27mu-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a4e74c522d630766b03a836c15bff77cb657c5fdf098abf8b1ada2aebc7d0819"},
|
||||
{file = "pycryptodome-3.21.0-cp27-cp27mu-musllinux_1_1_aarch64.whl", hash = "sha256:a3804675283f4764a02db05f5191eb8fec2bb6ca34d466167fc78a5f05bbe6b3"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-macosx_10_9_universal2.whl", hash = "sha256:2480ec2c72438430da9f601ebc12c518c093c13111a5c1644c82cdfc2e50b1e4"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:de18954104667f565e2fbb4783b56667f30fb49c4d79b346f52a29cb198d5b6b"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2de4b7263a33947ff440412339cb72b28a5a4c769b5c1ca19e33dd6cd1dcec6e"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0714206d467fc911042d01ea3a1847c847bc10884cf674c82e12915cfe1649f8"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7d85c1b613121ed3dbaa5a97369b3b757909531a959d229406a75b912dd51dd1"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:8898a66425a57bcf15e25fc19c12490b87bd939800f39a03ea2de2aea5e3611a"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-musllinux_1_2_i686.whl", hash = "sha256:932c905b71a56474bff8a9c014030bc3c882cee696b448af920399f730a650c2"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:18caa8cfbc676eaaf28613637a89980ad2fd96e00c564135bf90bc3f0b34dd93"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-win32.whl", hash = "sha256:280b67d20e33bb63171d55b1067f61fbd932e0b1ad976b3a184303a3dad22764"},
|
||||
{file = "pycryptodome-3.21.0-cp36-abi3-win_amd64.whl", hash = "sha256:b7aa25fc0baa5b1d95b7633af4f5f1838467f1815442b22487426f94e0d66c53"},
|
||||
{file = "pycryptodome-3.21.0-pp27-pypy_73-manylinux2010_x86_64.whl", hash = "sha256:2cb635b67011bc147c257e61ce864879ffe6d03342dc74b6045059dfbdedafca"},
|
||||
{file = "pycryptodome-3.21.0-pp27-pypy_73-win32.whl", hash = "sha256:4c26a2f0dc15f81ea3afa3b0c87b87e501f235d332b7f27e2225ecb80c0b1cdd"},
|
||||
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:d5ebe0763c982f069d3877832254f64974139f4f9655058452603ff559c482e8"},
|
||||
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ee86cbde706be13f2dec5a42b52b1c1d1cbb90c8e405c68d0755134735c8dc6"},
|
||||
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0fd54003ec3ce4e0f16c484a10bc5d8b9bd77fa662a12b85779a2d2d85d67ee0"},
|
||||
{file = "pycryptodome-3.21.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:5dfafca172933506773482b0e18f0cd766fd3920bd03ec85a283df90d8a17bc6"},
|
||||
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-macosx_10_15_x86_64.whl", hash = "sha256:590ef0898a4b0a15485b05210b4a1c9de8806d3ad3d47f74ab1dc07c67a6827f"},
|
||||
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f35e442630bc4bc2e1878482d6f59ea22e280d7121d7adeaedba58c23ab6386b"},
|
||||
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ff99f952db3db2fbe98a0b355175f93ec334ba3d01bbde25ad3a5a33abc02b58"},
|
||||
{file = "pycryptodome-3.21.0-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:8acd7d34af70ee63f9a849f957558e49a98f8f1634f86a59d2be62bb8e93f71c"},
|
||||
{file = "pycryptodome-3.21.0.tar.gz", hash = "sha256:f7787e0d469bdae763b876174cf2e6c0f7be79808af26b1da96f1a64bcf47297"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pydantic"
|
||||
version = "2.9.2"
|
||||
@@ -3509,4 +3551,4 @@ type = ["pytest-mypy"]
|
||||
[metadata]
|
||||
lock-version = "2.1"
|
||||
python-versions = ">=3.9.0,<4.0"
|
||||
content-hash = "eb85f0bcc0e8a715ef38afb58cf888f7c2ee8579ea6ed94900244365f24cddd9"
|
||||
content-hash = "b8641a0b2d92bee0363602e69f99b23366b2035b7e17ff017708194e6fbd0ac5"
|
||||
|
||||
@@ -37,6 +37,7 @@ uvloop = "0.21.0beta1"
|
||||
pyperf = "^2.7.0"
|
||||
py-spy = "^0.3.14"
|
||||
types-requests = "^2.32.0.20240914"
|
||||
pycryptodome = "^3.21.0"
|
||||
|
||||
[tool.ruff]
|
||||
lint.select = [ "E", "F", "I", "TID251" ]
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -377,6 +377,19 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge[sqlite_aes]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query --> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory]
|
||||
'''
|
||||
graph TD;
|
||||
@@ -797,6 +810,76 @@
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes].1
|
||||
dict({
|
||||
'definitions': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'inner': dict({
|
||||
'$ref': '#/definitions/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
]),
|
||||
'title': 'Input',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes].2
|
||||
dict({
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'title': 'Answer',
|
||||
'type': 'string',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'answer',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'Output',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory]
|
||||
'''
|
||||
graph TD;
|
||||
@@ -1217,6 +1300,76 @@
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
]),
|
||||
'title': 'Input',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes].2
|
||||
dict({
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'title': 'Answer',
|
||||
'type': 'string',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'answer',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'Output',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory]
|
||||
'''
|
||||
graph TD;
|
||||
@@ -1715,6 +1868,19 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite_aes]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_multiple_sinks_subgraphs
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
|
||||
@@ -16,6 +16,7 @@ from langgraph.checkpoint.postgres.aio import (
|
||||
AsyncPostgresSaver,
|
||||
AsyncShallowPostgresSaver,
|
||||
)
|
||||
from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
from langgraph.store.base import BaseStore
|
||||
@@ -61,6 +62,15 @@ def checkpointer_sqlite():
|
||||
yield checkpointer
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def checkpointer_sqlite_aes():
|
||||
with SqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes(
|
||||
key=b"1234567890123456"
|
||||
)
|
||||
yield checkpointer
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _checkpointer_sqlite_aio():
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
|
||||
@@ -437,6 +447,7 @@ REGULAR_CHECKPOINTERS_SYNC = [
|
||||
"postgres",
|
||||
"postgres_pipe",
|
||||
"postgres_pool",
|
||||
"sqlite_aes",
|
||||
]
|
||||
ALL_CHECKPOINTERS_SYNC = [
|
||||
*REGULAR_CHECKPOINTERS_SYNC,
|
||||
|
||||
Reference in New Issue
Block a user