github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-25 16:13:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
- **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
1 changed files with 16 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
libs/community/langchain_community/chat_message_histories/dynamodb.py
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) [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 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. 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__( def __init__(
@ -59,6 +61,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
ttl: Optional[int] = None, ttl: Optional[int] = None,
ttl_key_name: str = "expireAt", ttl_key_name: str = "expireAt",
history_size: Optional[int] = None, history_size: Optional[int] = None,
history_messages_key: Optional[str] = "History",
): ):
if boto3_session: if boto3_session:
client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url) client = boto3_session.resource("dynamodb", endpoint_url=endpoint_url)
@ -79,6 +82,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
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 self.history_size = history_size
self.history_messages_key = history_messages_key
if kms_key_id: if kms_key_id:
try: try:
@ -96,7 +100,9 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
actions = AttributeActions( actions = AttributeActions(
default_action=CryptoAction.DO_NOTHING, 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) aws_kms_cmp = AwsKmsCryptographicMaterialsProvider(key_id=kms_key_id)
self.table = EncryptedTable( self.table = EncryptedTable(
@ -126,7 +132,7 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
logger.error(error) logger.error(error)
if response and "Item" in response: if response and "Item" in response:
items = response["Item"]["History"] items = response["Item"][self.history_messages_key]
else: else:
items = [] items = []
@ -162,10 +168,16 @@ class DynamoDBChatMessageHistory(BaseChatMessageHistory):
expireAt = int(time.time()) + self.ttl expireAt = int(time.time()) + self.ttl
self.table.put_item( 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: 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: except ClientError as err:
logger.error(err) logger.error(err)