[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>
pull/11199/head
Justin Plock 12 months ago committed by GitHub
parent 4947ac2965
commit 926e4b6bad
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -38,6 +38,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
This argument is optional, but useful when using composite dynamodb keys, or This argument is optional, but useful when using composite dynamodb keys, or
isolating records based off of application details such as a user id. isolating records based off of application details such as a user id.
This may also contain global and local secondary index keys. 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__( def __init__(
@ -48,6 +50,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
primary_key_name: str = "SessionId", primary_key_name: str = "SessionId",
key: Optional[Dict[str, str]] = None, key: Optional[Dict[str, str]] = None,
boto3_session: Optional[Session] = None, boto3_session: Optional[Session] = None,
kms_key_id: Optional[str] = None,
): ):
if boto3_session: if boto3_session:
client = boto3_session.resource("dynamodb") client = boto3_session.resource("dynamodb")
@ -66,6 +69,32 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
self.session_id = session_id self.session_id = session_id
self.key: Dict = key or {primary_key_name: 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 @property
def messages(self) -> List[BaseMessage]: # type: ignore def messages(self) -> List[BaseMessage]: # type: ignore
"""Retrieve the messages from DynamoDB""" """Retrieve the messages from DynamoDB"""

Loading…
Cancel
Save