Skip to content

feat(storage): add content_filters to DynamoDbChatStorage #537

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nuthalapativarun wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(storage): add content_filters to DynamoDbChatStorage #537

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 76 additions & 20 deletions python/src/agent_squad/storage/dynamodb_chat_storage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import Union, Optional
import json
import time
import boto3
from agent_squad.storage import ChatStorage
Expand All @@ -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,
Expand All @@ -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} \
Expand All @@ -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)

Expand All @@ -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,
Expand All @@ -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):
Expand All @@ -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]

Expand All @@ -117,21 +168,23 @@ 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,
user_id: str,
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
Expand All @@ -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
Expand All @@ -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}"}]
Expand Down
105 changes: 105 additions & 0 deletions python/src/tests/storage/test_dynamodb_chat_storage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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):
"""
Expand Down