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

community[patch]: Implement TTL for DynamoDBChatMessageHistory (#15478)

Browse Source 
- **Description:** Implement TTL for DynamoDBChatMessageHistory, 
  - **Issue:** see #15477,
  - **Dependencies:** N/A,

---------

Co-authored-by: Piyush Jain <piyushjain@duck.com>
This commit is contained in:
Anthony Bernabeu 2024-01-29 19:22:46 +01:00 committed by GitHub
parent 1bc8d9a943
commit 2db79ab111
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 18 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
libs/community/langchain_community/chat_message_histories/dynamodb.py
View File

@ -38,6 +38,11 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
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 kms_key_id: an optional AWS KMS Key ID, AWS KMS Key ARN, or AWS KMS Alias for
client-side encryption client-side encryption
ttl: Optional Time-to-live (TTL) in seconds. Allows you to define a per-item
expiration timestamp that indicates when an item can be deleted from the
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)
""" """
def __init__( def __init__(
@ -49,6 +54,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
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, kms_key_id: Optional[str] = None,
ttl: Optional[int] = None,
ttl_key_name: str = "expireAt",
): ):
if boto3_session: if boto3_session:
client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url) client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url)
@ -66,6 +73,8 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
self.table = client.Table(table_name) self.table = client.Table(table_name)
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}
self.ttl = ttl
self.ttl_key_name = ttl_key_name
if kms_key_id: if kms_key_id:
try: try:
@ -134,7 +143,15 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
messages.append(_message) messages.append(_message)
try: try:
self.table.put_item(Item={**self.key, "History": messages}) if self.ttl:
import time
expireAt = int(time.time()) + self.ttl
self.table.put_item(
Item={**self.key, "History": messages, self.ttl_key_name: expireAt}
)
else:
self.table.put_item(Item={**self.key, "History": messages})
except ClientError as err: except ClientError as err:
logger.error(err) logger.error(err)