community[patch]: Dynamodb history messages key (#25658)

- **Description:** adding the history_messages_key so you don't have to
use "History" as a key in langchain
This commit is contained in:
Jakub W. 2024-08-23 10:05:28 +02:00 committed by GitHub
parent b28bc252c4
commit b865ee49a0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

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