From c17479a91393dc52e7d40eacbcb6927214036d0c Mon Sep 17 00:00:00 2001 From: Davi Schumacher Date: Sun, 26 May 2024 05:46:51 +0000 Subject: [PATCH 1/2] community: edit dynamodb chat message history to use update_item instead of put_item this allows someone to use the same chat session record to store additional attributes by updating the History attribute without overwriting the entire record --- .../chat_message_histories/dynamodb.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/libs/community/langchain_community/chat_message_histories/dynamodb.py b/libs/community/langchain_community/chat_message_histories/dynamodb.py index 7c0f5c3514226..ca1586a2b0a87 100644 --- a/libs/community/langchain_community/chat_message_histories/dynamodb.py +++ b/libs/community/langchain_community/chat_message_histories/dynamodb.py @@ -167,16 +167,19 @@ def add_message(self, message: BaseMessage) -> None: import time expireAt = int(time.time()) + self.ttl - self.table.put_item( - Item={ - **self.key, - self.history_messages_key: messages, - self.ttl_key_name: expireAt, - } + self.table.update_item( + Key={**self.key}, + UpdateExpression=( + f"set {self.history_messages_key} = :h, " + f"{self.ttl_key_name} = :t" + ), + ExpressionAttributeValues={":h": messages, ":t": expireAt}, ) else: - self.table.put_item( - Item={**self.key, self.history_messages_key: messages} + self.table.update_item( + Key={**self.key}, + UpdateExpression=f"set {self.history_messages_key} = :h", + ExpressionAttributeValues={":h": messages}, ) except ClientError as err: logger.error(err) From ebe148520cd7eee79999f8eb66e916a068f625d0 Mon Sep 17 00:00:00 2001 From: Davi Schumacher Date: Sun, 2 Jun 2024 01:33:18 +0000 Subject: [PATCH 2/2] community: add unit tests for the dynamodb chat message history class --- .../test_dynamodb_chat_message_history.py | 153 ++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 libs/community/tests/unit_tests/chat_message_histories/test_dynamodb_chat_message_history.py diff --git a/libs/community/tests/unit_tests/chat_message_histories/test_dynamodb_chat_message_history.py b/libs/community/tests/unit_tests/chat_message_histories/test_dynamodb_chat_message_history.py new file mode 100644 index 0000000000000..b88443e25fd02 --- /dev/null +++ b/libs/community/tests/unit_tests/chat_message_histories/test_dynamodb_chat_message_history.py @@ -0,0 +1,153 @@ +from typing import Any + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, messages_to_dict +from pytest_mock import MockerFixture + +from langchain_community.chat_message_histories.dynamodb import ( + DynamoDBChatMessageHistory, +) + +HISTORY_KEY = "ChatHistory" +TTL_KEY = "TimeToLive" + + +def dict_to_key(Key: dict) -> tuple: + return tuple(sorted(Key.items())) + + +class MockDynamoDBChatHistoryTable: + """Contains the table for the mock DynamoDB resource.""" + + class Table: + """Contains methods to mock Boto's DynamoDB calls""" + + def __init__(self, *args: tuple, **kwargs: dict[str, Any]) -> None: + self.items: dict = dict() + + def get_item(self, Key: dict) -> dict: + return self.items.get(dict_to_key(Key), dict()) + + def update_item( + self, Key: dict, UpdateExpression: str, ExpressionAttributeValues: dict + ) -> None: + update_dict = {HISTORY_KEY: ExpressionAttributeValues[":h"]} + + expression = UpdateExpression.split(", ") + if len(expression) > 1: + ttl_key_name = expression[1].replace(" = :t", "") + update_dict.update({ttl_key_name: ExpressionAttributeValues[":t"]}) + + self.items[dict_to_key(Key)] = {"Item": update_dict} + + def delete_item(self, Key: dict) -> None: + if dict_to_key(Key) in self.items.keys(): + del self.items[dict_to_key(Key)] + + +class MockBoto3DynamoDBSession: + """Creates a mock Boto session to return a DynamoDB table for testing + DynamoDBChatMessageHistory class methods.""" + + def resource( + self, *args: tuple, **kwargs: dict[str, Any] + ) -> MockDynamoDBChatHistoryTable: + return MockDynamoDBChatHistoryTable() + + +@pytest.fixture(scope="module") +def chat_history_config() -> dict: + return {"key": {"primaryKey": "foo", "secondaryKey": 123}, "ttl": 600} + + +@pytest.fixture(scope="class") +def ddb_chat_history_with_mock_boto_session( + chat_history_config: dict, +) -> DynamoDBChatMessageHistory: + return DynamoDBChatMessageHistory( + table_name="test_table", + session_id="test_session", + boto3_session=MockBoto3DynamoDBSession(), + key=chat_history_config["key"], + ttl=chat_history_config["ttl"], + ttl_key_name=TTL_KEY, + history_messages_key=HISTORY_KEY, + ) + + +class TestDynamoDBChatMessageHistory: + @pytest.mark.requires("botocore") + def test_add_message( + self, + mocker: MockerFixture, + ddb_chat_history_with_mock_boto_session: DynamoDBChatMessageHistory, + chat_history_config: dict, + ) -> None: + # For verifying the TTL value + mock_time_1 = 1234567000 + mock_time_2 = 1234568000 + + # Get the history class and mock DynamoDB table + history: DynamoDBChatMessageHistory = ddb_chat_history_with_mock_boto_session + history_table: MockDynamoDBChatHistoryTable.Table = history.table + history_item = history_table.get_item(chat_history_config["key"]) + assert history_item == dict() # Should be empty so far + + # Add the first message + mocker.patch("time.time", lambda: mock_time_1) + first_message = HumanMessage(content="new human message") + history.add_message(message=first_message) + item_after_human_message = history_table.get_item(chat_history_config["key"])[ + "Item" + ] + assert item_after_human_message[HISTORY_KEY] == messages_to_dict( + [first_message] + ) # History should only contain the first message + assert ( + item_after_human_message[TTL_KEY] + == mock_time_1 + chat_history_config["ttl"] + ) # TTL should exist + + # Add the second message + mocker.patch("time.time", lambda: mock_time_2) + second_message = AIMessage(content="new AI response") + history.add_message(message=second_message) + item_after_ai_message = history_table.get_item(chat_history_config["key"])[ + "Item" + ] + assert item_after_ai_message[HISTORY_KEY] == messages_to_dict( + [first_message, second_message] + ) # Second message should have appended + assert ( + item_after_ai_message[TTL_KEY] == mock_time_2 + chat_history_config["ttl"] + ) # TTL should have updated + + @pytest.mark.requires("botocore") + def test_clear( + self, + ddb_chat_history_with_mock_boto_session: DynamoDBChatMessageHistory, + chat_history_config: dict, + ) -> None: + # Get the history class and mock DynamoDB table + history: DynamoDBChatMessageHistory = ddb_chat_history_with_mock_boto_session + history_table: MockDynamoDBChatHistoryTable.Table = history.table + + # Use new key so we get a new chat session and add a message to the new session + new_session_key = {"primaryKey": "bar", "secondaryKey": 456} + history.key = new_session_key + history.add_message( + message=HumanMessage(content="human message for different chat session") + ) + + # Chat history table should now contain both chat sessions + assert set(history_table.items.keys()) == { + dict_to_key(chat_history_config["key"]), + dict_to_key(new_session_key), + } + + # Reset the key to the original and use the clear method + history.key = chat_history_config["key"] + history.clear() + + # Only the original chat session should be removed + assert set(history_table.items.keys()) == {dict_to_key(new_session_key)}