community[patch]: history size support for DynamoDBChatMessageHistory (#16794)

**Description:** PR adds support for limiting number of messages
preserved in a session history for DynamoDBChatMessageHistory

---------

Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>
Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
Luka Krapic 2024-03-29 19:56:21 +01:00 committed by GitHub
parent 6dbf1a2de0
commit 727a2ea9f1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -43,6 +43,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
table. DynamoDB handles deletion of expired items without consuming table. DynamoDB handles deletion of expired items without consuming
write throughput. To enable this feature on the table, follow the write throughput. To enable this feature on the table, follow the
[AWS DynamoDB documentation](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/time-to-live-ttl-how-to.html) [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.
""" """
def __init__( def __init__(
@ -56,6 +58,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
kms_key_id: Optional[str] = None, kms_key_id: Optional[str] = None,
ttl: Optional[int] = None, ttl: Optional[int] = None,
ttl_key_name: str = "expireAt", ttl_key_name: str = "expireAt",
history_size: Optional[int] = None,
): ):
if boto3_session: if boto3_session:
client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url) client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url)
@ -75,6 +78,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
self.key: Dict = key or {primary_key_name: session_id} self.key: Dict = key or {primary_key_name: session_id}
self.ttl = ttl self.ttl = ttl
self.ttl_key_name = ttl_key_name self.ttl_key_name = ttl_key_name
self.history_size = history_size
if kms_key_id: if kms_key_id:
try: try:
@ -149,6 +153,9 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
_message = message_to_dict(message) _message = message_to_dict(message)
messages.append(_message) messages.append(_message)
if self.history_size:
messages = messages[-self.history_size :]
try: try:
if self.ttl: if self.ttl:
import time import time