mirror of
https://github.com/hwchase17/langchain.git
synced 2025-06-01 04:29:09 +00:00
[Feat] Add optional client-side encryption to DynamoDB chat history memory (#11115)
**Description:** Added optional client-side encryption to the Amazon DynamoDB chat history memory with an AWS KMS Key ID using the [AWS Database Encryption SDK for Python](https://docs.aws.amazon.com/database-encryption-sdk/latest/devguide/python.html) **Issue:** #7886 **Dependencies:** [dynamodb-encryption-sdk](https://pypi.org/project/dynamodb-encryption-sdk/) **Tag maintainer:** @hwchase17 **Twitter handle:** [@jplock](https://twitter.com/jplock/) --------- Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
parent
4947ac2965
commit
926e4b6bad
@ -38,6 +38,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
|
||||
This argument is optional, but useful when using composite dynamodb keys, or
|
||||
isolating records based off of application details such as a user id.
|
||||
This may also contain global and local secondary index keys.
|
||||
kms_key_id: an optional AWS KMS Key ID, AWS KMS Key ARN, or AWS KMS Alias for
|
||||
client-side encryption
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@ -48,6 +50,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
|
||||
primary_key_name: str = "SessionId",
|
||||
key: Optional[Dict[str, str]] = None,
|
||||
boto3_session: Optional[Session] = None,
|
||||
kms_key_id: Optional[str] = None,
|
||||
):
|
||||
if boto3_session:
|
||||
client = boto3_session.resource("dynamodb")
|
||||
@ -66,6 +69,32 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
|
||||
self.session_id = session_id
|
||||
self.key: Dict = key or {primary_key_name: session_id}
|
||||
|
||||
if kms_key_id:
|
||||
try:
|
||||
from dynamodb_encryption_sdk.encrypted.table import EncryptedTable
|
||||
from dynamodb_encryption_sdk.identifiers import CryptoAction
|
||||
from dynamodb_encryption_sdk.material_providers.aws_kms import (
|
||||
AwsKmsCryptographicMaterialsProvider,
|
||||
)
|
||||
from dynamodb_encryption_sdk.structures import AttributeActions
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Unable to import dynamodb_encryption_sdk, please install with "
|
||||
"`pip install dynamodb-encryption-sdk`."
|
||||
) from e
|
||||
|
||||
actions = AttributeActions(
|
||||
default_action=CryptoAction.DO_NOTHING,
|
||||
attribute_actions={"History": CryptoAction.ENCRYPT_AND_SIGN},
|
||||
)
|
||||
aws_kms_cmp = AwsKmsCryptographicMaterialsProvider(key_id=kms_key_id)
|
||||
self.table = EncryptedTable(
|
||||
table=self.table,
|
||||
materials_provider=aws_kms_cmp,
|
||||
attribute_actions=actions,
|
||||
auto_refresh_table_indexes=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def messages(self) -> List[BaseMessage]: # type: ignore
|
||||
"""Retrieve the messages from DynamoDB"""
|
||||
|
Loading…
Reference in New Issue
Block a user