From 0adaf8235f2d5721d6d5d236c028500b58a0d3a0 Mon Sep 17 00:00:00 2001 From: Varun Nuthalapati Date: Sun, 28 Jun 2026 12:33:32 -0700 Subject: [PATCH] feat(storage): add content_filters to DynamoDbChatStorage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an optional content_filters parameter — a list of (original, placeholder) pairs — to DynamoDbChatStorage. Sensitive values are replaced with their placeholders before writing to DynamoDB and restored on every read path (fetch_chat, fetch_chat_with_timestamp, fetch_all_chats). Closes #131 --- .../storage/dynamodb_chat_storage.py | 96 ++++++++++++---- .../storage/test_dynamodb_chat_storage.py | 105 ++++++++++++++++++ 2 files changed, 181 insertions(+), 20 deletions(-) diff --git a/python/src/agent_squad/storage/dynamodb_chat_storage.py b/python/src/agent_squad/storage/dynamodb_chat_storage.py index 7c327c056..d312a522d 100644 --- a/python/src/agent_squad/storage/dynamodb_chat_storage.py +++ b/python/src/agent_squad/storage/dynamodb_chat_storage.py @@ -1,4 +1,5 @@ from typing import Union, Optional +import json import time import boto3 from agent_squad.storage import ChatStorage @@ -13,15 +14,49 @@ def __init__(self, table_name: str, region: str, ttl_key: Optional[str] = None, - ttl_duration: int = 3600): + ttl_duration: int = 3600, + content_filters: Optional[list[tuple[str, str]]] = None): super().__init__() self.table_name = table_name self.ttl_key = ttl_key self.ttl_duration = int(ttl_duration) + self.content_filters = content_filters or [] self.dynamodb = boto3.resource('dynamodb', region_name=region) self.table = self.dynamodb.Table(table_name) user_agent.register_feature_to_resource(self.dynamodb, feature='storage-ddb') + def _apply_content_filters(self, content: list) -> list: + if not self.content_filters: + return content + serialized = json.dumps(content) + for original, placeholder in self.content_filters: + serialized = serialized.replace(original, placeholder) + return json.loads(serialized) + + def _reverse_content_filters(self, content: list) -> list: + if not self.content_filters: + return content + serialized = json.dumps(content) + for original, placeholder in self.content_filters: + serialized = serialized.replace(placeholder, original) + return json.loads(serialized) + + async def _fetch_raw( + self, + user_id: str, + session_id: str, + agent_id: str + ) -> list[TimestampedMessage]: + key = self._generate_key(user_id, session_id, agent_id) + try: + response = self.table.get_item(Key={'PK': user_id, 'SK': key}) + return self._dict_to_conversation( + response.get('Item', {}).get('conversation', []) + ) + except Exception as error: + Logger.error(f"Error getting conversation from DynamoDB: {str(error)}") + raise error + async def save_chat_message( self, user_id: str, @@ -31,7 +66,7 @@ async def save_chat_message( max_history_size: Optional[int] = None ) -> list[ConversationMessage]: key = self._generate_key(user_id, session_id, agent_id) - existing_conversation = await self.fetch_chat_with_timestamp(user_id, session_id, agent_id) + existing_conversation = await self._fetch_raw(user_id, session_id, agent_id) if self.is_same_role_as_last_message(existing_conversation, new_message): Logger.debug(f"> Consecutive {new_message.role} \ @@ -41,7 +76,12 @@ async def save_chat_message( if isinstance(new_message, ConversationMessage): new_message = TimestampedMessage( role=new_message.role, - content=new_message.content) + content=self._apply_content_filters(new_message.content)) + elif self.content_filters: + new_message = TimestampedMessage( + role=new_message.role, + content=self._apply_content_filters(new_message.content), + timestamp=new_message.timestamp) existing_conversation.append(new_message) @@ -65,7 +105,10 @@ async def save_chat_message( Logger.error(f"Error saving conversation to DynamoDB:{str(error)}") raise error - return self._remove_timestamps(trimmed_conversation) + return [ConversationMessage( + role=msg.role, + content=self._reverse_content_filters(msg.content) + ) for msg in trimmed_conversation] async def save_chat_messages(self, user_id: str, @@ -79,7 +122,7 @@ async def save_chat_messages(self, Save multiple messages at once """ key = self._generate_key(user_id, session_id, agent_id) - existing_conversation = await self.fetch_chat_with_timestamp(user_id, session_id, agent_id) + existing_conversation = await self._fetch_raw(user_id, session_id, agent_id) #TODO: check messages are consecutive # if self.is_same_role_as_last_message(existing_conversation, new_messages): @@ -91,7 +134,15 @@ async def save_chat_messages(self, new_messages = [ TimestampedMessage( role=new_message.role, - content=new_message.content + content=self._apply_content_filters(new_message.content) + ) + for new_message in new_messages] + elif self.content_filters: + new_messages = [ + TimestampedMessage( + role=new_message.role, + content=self._apply_content_filters(new_message.content), + timestamp=new_message.timestamp ) for new_message in new_messages] @@ -117,7 +168,10 @@ async def save_chat_messages(self, Logger.error(f"Error saving conversation to DynamoDB:{str(error)}") raise error - return self._remove_timestamps(trimmed_conversation) + return [ConversationMessage( + role=msg.role, + content=self._reverse_content_filters(msg.content) + ) for msg in trimmed_conversation] async def fetch_chat( self, @@ -125,13 +179,12 @@ async def fetch_chat( session_id: str, agent_id: str ) -> list[ConversationMessage]: - key = self._generate_key(user_id, session_id, agent_id) try: - response = self.table.get_item(Key={'PK': user_id, 'SK': key}) - stored_messages: list[TimestampedMessage] = self._dict_to_conversation( - response.get('Item', {}).get('conversation', []) - ) - return self._remove_timestamps(stored_messages) + stored_messages = await self._fetch_raw(user_id, session_id, agent_id) + return [ConversationMessage( + role=msg.role, + content=self._reverse_content_filters(msg.content) + ) for msg in stored_messages] except Exception as error: Logger.error(f"Error getting conversation from DynamoDB:{str(error)}") raise error @@ -142,13 +195,15 @@ async def fetch_chat_with_timestamp( session_id: str, agent_id: str ) -> list[TimestampedMessage]: - key = self._generate_key(user_id, session_id, agent_id) try: - response = self.table.get_item(Key={'PK': user_id, 'SK': key}) - stored_messages: list[TimestampedMessage] = self._dict_to_conversation( - response.get('Item', {}).get('conversation', []) - ) - return stored_messages + stored_messages = await self._fetch_raw(user_id, session_id, agent_id) + if not self.content_filters: + return stored_messages + return [TimestampedMessage( + role=msg.role, + content=self._reverse_content_filters(msg.content), + timestamp=msg.timestamp + ) for msg in stored_messages] except Exception as error: Logger.error(f"Error getting conversation from DynamoDB: {str(error)}") raise error @@ -174,7 +229,8 @@ async def fetch_all_chats(self, user_id: str, session_id: str) -> list[Conversat agent_id = item['SK'].split('#')[1] for msg in item['conversation']: - content = msg['content'] + content = self._reverse_content_filters(msg['content']) \ + if isinstance(msg['content'], list) else msg['content'] if msg['role'] == ParticipantRole.ASSISTANT.value: text = content[0]['text'] if isinstance(content, list) else content content = [{'text': f"[{agent_id}] {text}"}] diff --git a/python/src/tests/storage/test_dynamodb_chat_storage.py b/python/src/tests/storage/test_dynamodb_chat_storage.py index 1893bffcd..fb3d2f174 100644 --- a/python/src/tests/storage/test_dynamodb_chat_storage.py +++ b/python/src/tests/storage/test_dynamodb_chat_storage.py @@ -27,6 +27,14 @@ def dynamodb_table(): def chat_storage(dynamodb_table): return DynamoDbChatStorage(table_name='test_table', region='us-east-1', ttl_key='TTL', ttl_duration=3600) +@pytest.fixture +def filtered_chat_storage(dynamodb_table): + return DynamoDbChatStorage( + table_name='test_table', + region='us-east-1', + content_filters=[('acme-corp', '#COMPANY#'), ('secret-token-xyz', '#TOKEN#')] + ) + @pytest.mark.asyncio async def test_save_and_fetch_chat_message(chat_storage): user_id = 'user1' @@ -144,6 +152,103 @@ async def test_save_and_fetch_chat_messages(chat_storage): assert fetched_messages[4].role == ParticipantRole.USER.value +@pytest.mark.asyncio +async def test_content_filters_replace_on_save_and_restore_on_fetch(filtered_chat_storage): + """Sensitive values should be stored as placeholders and restored on fetch.""" + user_id = 'user1' + session_id = 'session1' + agent_id = 'agent1' + message = ConversationMessage( + role=ParticipantRole.USER.value, + content=[{'text': 'Hello from acme-corp'}] + ) + + await filtered_chat_storage.save_chat_message(user_id, session_id, agent_id, message) + + # Verify placeholder is stored in DynamoDB + raw = filtered_chat_storage.table.get_item( + Key={'PK': user_id, 'SK': filtered_chat_storage._generate_key(user_id, session_id, agent_id)} + ) + stored_text = raw['Item']['conversation'][0]['content'][0]['text'] + assert stored_text == 'Hello from #COMPANY#' + + # Verify original value is returned by fetch_chat + fetched = await filtered_chat_storage.fetch_chat(user_id, session_id, agent_id) + assert fetched[0].content == [{'text': 'Hello from acme-corp'}] + + +@pytest.mark.asyncio +async def test_content_filters_fetch_chat_with_timestamp(filtered_chat_storage): + """fetch_chat_with_timestamp must also restore original values.""" + user_id = 'user1' + session_id = 'session1' + agent_id = 'agent1' + message = ConversationMessage( + role=ParticipantRole.USER.value, + content=[{'text': 'token: secret-token-xyz'}] + ) + + await filtered_chat_storage.save_chat_message(user_id, session_id, agent_id, message) + + fetched = await filtered_chat_storage.fetch_chat_with_timestamp(user_id, session_id, agent_id) + assert fetched[0].content == [{'text': 'token: secret-token-xyz'}] + assert isinstance(fetched[0].timestamp, Decimal) + + +@pytest.mark.asyncio +async def test_content_filters_multiple_filters_round_trip(filtered_chat_storage): + """Multiple filters should all be applied and reversed correctly.""" + user_id = 'user1' + session_id = 'session1' + agent_id = 'agent1' + + msg1 = ConversationMessage( + role=ParticipantRole.USER.value, + content=[{'text': 'Company acme-corp uses secret-token-xyz'}] + ) + msg2 = ConversationMessage( + role=ParticipantRole.ASSISTANT.value, + content=[{'text': 'Confirmed for acme-corp'}] + ) + + await filtered_chat_storage.save_chat_message(user_id, session_id, agent_id, msg1) + await filtered_chat_storage.save_chat_message(user_id, session_id, agent_id, msg2) + + fetched = await filtered_chat_storage.fetch_chat(user_id, session_id, agent_id) + assert fetched[0].content == [{'text': 'Company acme-corp uses secret-token-xyz'}] + assert fetched[1].content == [{'text': 'Confirmed for acme-corp'}] + + +@pytest.mark.asyncio +async def test_content_filters_save_chat_messages(filtered_chat_storage): + """save_chat_messages should also apply content filters.""" + messages = [ + ConversationMessage(role=ParticipantRole.USER.value, content=[{'text': 'acme-corp question'}]), + ConversationMessage(role=ParticipantRole.ASSISTANT.value, content=[{'text': 'acme-corp answer'}]), + ] + + await filtered_chat_storage.save_chat_messages('user1', 'session1', 'agent1', messages) + fetched = await filtered_chat_storage.fetch_chat('user1', 'session1', 'agent1') + + assert fetched[0].content == [{'text': 'acme-corp question'}] + assert fetched[1].content == [{'text': 'acme-corp answer'}] + + +@pytest.mark.asyncio +async def test_no_content_filters_behaviour_unchanged(chat_storage): + """Storage without content_filters must behave identically to the original.""" + user_id = 'user1' + session_id = 'session1' + agent_id = 'agent1' + message = ConversationMessage(role=ParticipantRole.USER.value, content=[{'text': 'plain text'}]) + + saved = await chat_storage.save_chat_message(user_id, session_id, agent_id, message) + assert saved[0].content == [{'text': 'plain text'}] + + fetched = await chat_storage.fetch_chat(user_id, session_id, agent_id) + assert fetched[0].content == [{'text': 'plain text'}] + + @pytest.mark.asyncio async def test_save_and_fetch_chat_messages_timestamp(chat_storage): """