diff --git a/libs/community/langchain_community/chat_message_histories/dynamodb.py b/libs/community/langchain_community/chat_message_histories/dynamodb.py index 2f1c18c45c..630781310e 100644 --- a/libs/community/langchain_community/chat_message_histories/dynamodb.py +++ b/libs/community/langchain_community/chat_message_histories/dynamodb.py @@ -43,6 +43,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): table. DynamoDB handles deletion of expired items without consuming 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) + 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__( @@ -56,6 +58,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): kms_key_id: Optional[str] = None, ttl: Optional[int] = None, ttl_key_name: str = "expireAt", + history_size: Optional[int] = None, ): if boto3_session: 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.ttl = ttl self.ttl_key_name = ttl_key_name + self.history_size = history_size if kms_key_id: try: @@ -149,6 +153,9 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory): _message = message_to_dict(message) messages.append(_message) + if self.history_size: + messages = messages[-self.history_size :] + try: if self.ttl: import time