Skip to content

summarization manager - add summary prompt to messages #698

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.

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
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ def __init__(
self.summarization_agent = summarization_agent
self.summarization_system_prompt = summarization_system_prompt
self._summary_message: Optional[Message] = None
self._summary_prompt: Message = {"content": [{"text": "Please summarize this conversation"}], "role": "user"}

@override
def restore_from_session(self, state: dict[str, Any]) -> Optional[list[Message]]:
Expand All @@ -94,7 +95,7 @@ def restore_from_session(self, state: dict[str, Any]) -> Optional[list[Message]]
"""
super().restore_from_session(state)
self._summary_message = state.get("summary_message")
return [self._summary_message] if self._summary_message else None
return [self._summary_prompt, self._summary_message] if self._summary_message else None
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the one change I am somewhat worried about. The return type is Optional[list[Message]] and so I feel it should be okay to include the summary prompt. Here is an example of how the method is used.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ohh this is clever! At first I didnt think this would work since we only store a single message in the conversation manager state (ref), but since this message is static it should be fine!


def get_state(self) -> dict[str, Any]:
"""Returns a dictionary representation of the state for the Summarizing Conversation Manager."""
Expand Down Expand Up @@ -152,15 +153,15 @@ def reduce_context(self, agent: "Agent", e: Optional[Exception] = None, **kwargs

# Keep track of the number of messages that have been summarized thus far.
self.removed_message_count += len(messages_to_summarize)
# If there is a summary message, don't count it in the removed_message_count.
# If there is a summary message, don't count it or the summary prompt in the removed_message_count.
if self._summary_message:
self.removed_message_count -= 1
self.removed_message_count -= 2

# Generate summary
self._summary_message = self._generate_summary(messages_to_summarize, agent)

# Replace the summarized messages with the summary
agent.messages[:] = [self._summary_message] + remaining_messages
agent.messages[:] = [self._summary_prompt, self._summary_message] + remaining_messages

except Exception as summarization_error:
logger.error("Summarization failed: %s", summarization_error)
Expand Down Expand Up @@ -200,8 +201,7 @@ def _generate_summary(self, messages: List[Message], agent: "Agent") -> Message:
summarization_agent.messages = messages

# Use the agent to generate summary with rich content (can use tools if needed)
result = summarization_agent("Please summarize this conversation.")

result = summarization_agent(self._summary_prompt["content"])
return result.message

finally:
Expand Down
54 changes: 28 additions & 26 deletions tests/strands/agent/test_summarizing_conversation_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,13 +95,16 @@ def test_reduce_context_with_summarization(summarizing_manager, mock_agent):

summarizing_manager.reduce_context(mock_agent)

# Should have: 1 summary message + 2 preserved recent messages + remaining from summarization
assert len(mock_agent.messages) == 4
# Should have: 1 summary prompt + 1 summary message + 2 preserved recent messages + remaining from summarization
assert len(mock_agent.messages) == 5

# First message should be the summary
assert mock_agent.messages[0]["role"] == "assistant"
first_content = mock_agent.messages[0]["content"][0]
assert "text" in first_content and "This is a summary of the conversation." in first_content["text"]
# First message should be the summary prompt
assert mock_agent.messages[0] == {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
# Second message should be the summary
assert mock_agent.messages[1] == {
"content": [{"text": "This is a summary of the conversation."}],
"role": "assistant",
}

# Recent messages should be preserved
assert "Message 3" in str(mock_agent.messages[-2]["content"])
Expand Down Expand Up @@ -434,17 +437,18 @@ def test_reduce_context_tool_pair_adjustment_works_with_forward_search():
# messages_to_summarize_count = (3 - 1) * 0.5 = 1
# But split point adjustment will move forward from the toolUse, potentially increasing count
manager.reduce_context(mock_agent)
# Should have summary + remaining messages
assert len(mock_agent.messages) == 2

# First message should be the summary
assert mock_agent.messages[0]["role"] == "assistant"
summary_content = mock_agent.messages[0]["content"][0]
assert "text" in summary_content and "This is a summary of the conversation." in summary_content["text"]

# Should have summary prompt + summary + remaining messages
assert len(mock_agent.messages) == 3

# First message should be the summary prompt
assert mock_agent.messages[0] == {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
# Second message should be the summary
assert mock_agent.messages[1] == {
"content": [{"text": "This is a summary of the conversation."}],
"role": "assistant",
}
# Last message should be the preserved recent message
assert mock_agent.messages[1]["role"] == "user"
assert mock_agent.messages[1]["content"][0]["text"] == "Latest message"
assert mock_agent.messages[2] == {"content": [{"text": "Latest message"}], "role": "user"}


def test_adjust_split_point_exceeds_message_length(summarizing_manager):
Expand Down Expand Up @@ -590,21 +594,19 @@ def test_summarizing_conversation_manager_properly_records_removed_message_count
assert manager.removed_message_count == 0

manager.reduce_context(agent)
# Assert the oldest message is the sumamry message
assert manager._summary_message["content"][0]["text"] == "Summary"
# There are 8 messages in the agent messages array, since half will be summarized,
# 4 will remain plus 1 summary message = 5
assert (len(agent.messages)) == 5
# 4 will remain plus 1 summary prompt and 1 summary message = 6
assert len(agent.messages) == 6
# Half of the messages were summarized and removed: 8/2 = 4
assert manager.removed_message_count == 4

manager.reduce_context(agent)
assert manager._summary_message["content"][0]["text"] == "Summary"
# After the first summary, 5 messages remain. Summarizing again will lead to:
# 5 - (int(5/2)) (messages to be sumamrized) + 1 (new summary message) = 5 - 2 + 1 = 4
assert (len(agent.messages)) == 4
# Half of the messages were summarized and removed: int(5/2) = 2
# However, one of the messages that was summarized was the previous summary message,
# so we dont count this toward the total:
# 4 (Previously removed messages) + 2 (removed messages) - 1 (Previous summary message) = 5
# After the first summary, 6 messages remain. Summarizing again will lead to:
# 6 - (6/2) (messages to be sumamrized) + 1 (summary prompt) + 1 (new summary message) = 6 - 3 + 1 + 1 = 5
assert len(agent.messages) == 5
# Half of the messages were summarized and removed: (6/2) = 3
# However, the summary prompt and previous summary were also summarized but we don't count this in the total:
# 4 (Previously removed messages) + 3 (removed messages) - 1 (summary prompt) - 1 (Previous summary message) = 5
assert manager.removed_message_count == 5
7 changes: 3 additions & 4 deletions tests/strands/session/test_repository_session_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,10 +151,9 @@ def test_initialize_restores_existing_agent_with_summarizing_conversation_manage

# Verify agent state restored
assert agent.state.get("key") == "value"
# The session message plus the summary message
assert len(agent.messages) == 2
assert agent.messages[1]["role"] == "user"
assert agent.messages[1]["content"][0]["text"] == "Hello"
# The session message plus the summary prompt plus the summary message
assert len(agent.messages) == 3
assert agent.messages[2] == {"content": [{"text": "Hello"}], "role": "user"}
assert agent.conversation_manager.removed_message_count == 1


Expand Down
15 changes: 9 additions & 6 deletions tests_integ/test_summarizing_conversation_manager_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,15 +151,18 @@ def test_summarization_with_context_overflow(model):

# Verify summarization occurred
assert len(agent.messages) < initial_message_count
# Should have: 1 summary + remaining messages
# Should have: 1 summary prompt + 1 summary + remaining messages
# With 6 messages, summary_ratio=0.5, preserve_recent_messages=2:
# messages_to_summarize = min(6 * 0.5, 6 - 2) = min(3, 4) = 3
# So we summarize 3 messages, leaving 3 remaining + 1 summary = 4 total
expected_total_messages = 4
# So we summarize 3 messages, leaving 3 remaining + 1 summary prompt + 1 summary = 5 total
expected_total_messages = 5
assert len(agent.messages) == expected_total_messages

# First message should be the summary (assistant message)
summary_message = agent.messages[0]
# First message should be the summary prompt
summary_prompt = agent.messages[0]
assert summary_prompt == {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
# Second message should be the summary
summary_message = agent.messages[1]
assert summary_message["role"] == "assistant"
assert len(summary_message["content"]) > 0

Expand Down Expand Up @@ -361,7 +364,7 @@ def test_dedicated_summarization_agent(model, summarization_model):
assert len(agent.messages) < original_length

# Get the summary message
summary_message = agent.messages[0]
summary_message = agent.messages[1]
assert summary_message["role"] == "assistant"

# Extract summary text
Expand Down
Loading