diff --git a/libs/community/langchain_community/chat_message_histories/dynamodb.py b/libs/community/langchain_community/chat_message_histories/dynamodb.py index 630781310e9..7c0f5c35142 100644 --- a/libs/community/langchain_community/chat_message_histories/dynamodb.py +++ b/libs/community/langchain_community/chat_message_histories/dynamodb.py @@ -45,6 +45,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): [AWS DynamoDB documentation](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/time-to-live-ttl-how-to.html) history_size: Maximum number of messages to store. If None then there is no limit. If not None then only the latest `history_size` messages are stored. + history_messages_key: Key for the chat history where the messages + are stored and updated """ def __init__( @@ -59,6 +61,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): ttl: Optional[int] = None, ttl_key_name: str = "expireAt", history_size: Optional[int] = None, + history_messages_key: Optional[str] = "History", ): if boto3_session: client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url) @@ -79,6 +82,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): self.ttl = ttl self.ttl_key_name = ttl_key_name self.history_size = history_size + self.history_messages_key = history_messages_key if kms_key_id: try: @@ -96,7 +100,9 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): actions = AttributeActions( default_action=CryptoAction.DO_NOTHING, - attribute_actions={"History": CryptoAction.ENCRYPT_AND_SIGN}, + attribute_actions={ + self.history_messages_key: CryptoAction.ENCRYPT_AND_SIGN + }, ) aws_kms_cmp = AwsKmsCryptographicMaterialsProvider(key_id=kms_key_id) self.table = EncryptedTable( @@ -126,7 +132,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): logger.error(error) if response and "Item" in response: - items = response["Item"]["History"] + items = response["Item"][self.history_messages_key] else: items = [] @@ -162,10 +168,16 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): expireAt = int(time.time()) + self.ttl self.table.put_item( - Item={**self.key, "History": messages, self.ttl_key_name: expireAt} + Item={ + **self.key, + self.history_messages_key: messages, + self.ttl_key_name: expireAt, + } ) else: - self.table.put_item(Item={**self.key, "History": messages}) + self.table.put_item( + Item={**self.key, self.history_messages_key: messages} + ) except ClientError as err: logger.error(err)