github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-21 03:51:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Add optional client-side encryption to DynamoDB chat history memory (#11115)

Browse Source 
**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:
Justin Plock 2023-09-28 16:29:46 -04:00 committed by GitHub
parent 4947ac2965
commit 926e4b6bad
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 29 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
libs/langchain/langchain/memory/chat_message_histories/dynamodb.py
View File

@ -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"""