Skip to content

Commit 39c0665

Browse files
committed
summarization manager - add summary prompt to messages
1 parent 17ccdd2 commit 39c0665

File tree

4 files changed

+46
-42
lines changed

4 files changed

+46
-42
lines changed

src/strands/agent/conversation_manager/summarizing_conversation_manager.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ def __init__(
8181
self.summarization_agent = summarization_agent
8282
self.summarization_system_prompt = summarization_system_prompt
8383
self._summary_message: Optional[Message] = None
84+
self._summary_prompt: Message = {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
8485

8586
@override
8687
def restore_from_session(self, state: dict[str, Any]) -> Optional[list[Message]]:
@@ -94,7 +95,7 @@ def restore_from_session(self, state: dict[str, Any]) -> Optional[list[Message]]
9495
"""
9596
super().restore_from_session(state)
9697
self._summary_message = state.get("summary_message")
97-
return [self._summary_message] if self._summary_message else None
98+
return [self._summary_prompt, self._summary_message] if self._summary_message else None
9899

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

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

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

162163
# Replace the summarized messages with the summary
163-
agent.messages[:] = [self._summary_message] + remaining_messages
164+
agent.messages[:] = [self._summary_prompt, self._summary_message] + remaining_messages
164165

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

202203
# Use the agent to generate summary with rich content (can use tools if needed)
203-
result = summarization_agent("Please summarize this conversation.")
204-
204+
result = summarization_agent(self._summary_prompt["content"])
205205
return result.message
206206

207207
finally:

tests/strands/agent/test_summarizing_conversation_manager.py

Lines changed: 28 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -95,13 +95,16 @@ def test_reduce_context_with_summarization(summarizing_manager, mock_agent):
9595

9696
summarizing_manager.reduce_context(mock_agent)
9797

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

101-
# First message should be the summary
102-
assert mock_agent.messages[0]["role"] == "assistant"
103-
first_content = mock_agent.messages[0]["content"][0]
104-
assert "text" in first_content and "This is a summary of the conversation." in first_content["text"]
101+
# First message should be the summary prompt
102+
assert mock_agent.messages[0] == {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
103+
# Second message should be the summary
104+
assert mock_agent.messages[1] == {
105+
"content": [{"text": "This is a summary of the conversation."}],
106+
"role": "assistant",
107+
}
105108

106109
# Recent messages should be preserved
107110
assert "Message 3" in str(mock_agent.messages[-2]["content"])
@@ -434,17 +437,18 @@ def test_reduce_context_tool_pair_adjustment_works_with_forward_search():
434437
# messages_to_summarize_count = (3 - 1) * 0.5 = 1
435438
# But split point adjustment will move forward from the toolUse, potentially increasing count
436439
manager.reduce_context(mock_agent)
437-
# Should have summary + remaining messages
438-
assert len(mock_agent.messages) == 2
439-
440-
# First message should be the summary
441-
assert mock_agent.messages[0]["role"] == "assistant"
442-
summary_content = mock_agent.messages[0]["content"][0]
443-
assert "text" in summary_content and "This is a summary of the conversation." in summary_content["text"]
444-
440+
# Should have summary prompt + summary + remaining messages
441+
assert len(mock_agent.messages) == 3
442+
443+
# First message should be the summary prompt
444+
assert mock_agent.messages[0] == {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
445+
# Second message should be the summary
446+
assert mock_agent.messages[1] == {
447+
"content": [{"text": "This is a summary of the conversation."}],
448+
"role": "assistant",
449+
}
445450
# Last message should be the preserved recent message
446-
assert mock_agent.messages[1]["role"] == "user"
447-
assert mock_agent.messages[1]["content"][0]["text"] == "Latest message"
451+
assert mock_agent.messages[2] == {"content": [{"text": "Latest message"}], "role": "user"}
448452

449453

450454
def test_adjust_split_point_exceeds_message_length(summarizing_manager):
@@ -590,21 +594,19 @@ def test_summarizing_conversation_manager_properly_records_removed_message_count
590594
assert manager.removed_message_count == 0
591595

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

601604
manager.reduce_context(agent)
602605
assert manager._summary_message["content"][0]["text"] == "Summary"
603-
# After the first summary, 5 messages remain. Summarizing again will lead to:
604-
# 5 - (int(5/2)) (messages to be sumamrized) + 1 (new summary message) = 5 - 2 + 1 = 4
605-
assert (len(agent.messages)) == 4
606-
# Half of the messages were summarized and removed: int(5/2) = 2
607-
# However, one of the messages that was summarized was the previous summary message,
608-
# so we dont count this toward the total:
609-
# 4 (Previously removed messages) + 2 (removed messages) - 1 (Previous summary message) = 5
606+
# After the first summary, 6 messages remain. Summarizing again will lead to:
607+
# 6 - (6/2) (messages to be sumamrized) + 1 (summary prompt) + 1 (new summary message) = 6 - 3 + 1 + 1 = 5
608+
assert len(agent.messages) == 5
609+
# Half of the messages were summarized and removed: (6/2) = 3
610+
# However, the summary prompt and previous summary message was also summarized but we don't count this in the total.
611+
# 4 (Previously removed messages) + 3 (removed messages) - 1 (summary prompt) - 1 (Previous summary message) = 5
610612
assert manager.removed_message_count == 5

tests/strands/session/test_repository_session_manager.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -151,10 +151,9 @@ def test_initialize_restores_existing_agent_with_summarizing_conversation_manage
151151

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

160159

tests_integ/test_summarizing_conversation_manager_integration.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -151,15 +151,18 @@ def test_summarization_with_context_overflow(model):
151151

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

161-
# First message should be the summary (assistant message)
162-
summary_message = agent.messages[0]
161+
# First message should be the summary prompt
162+
summary_prompt = agent.messages[0]
163+
assert summary_prompt == {"content": [{"text": "Please summarize this conversation"}], "role": "user"}
164+
# Second message should be the summary
165+
summary_message = agent.messages[1]
163166
assert summary_message["role"] == "assistant"
164167
assert len(summary_message["content"]) > 0
165168

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

363366
# Get the summary message
364-
summary_message = agent.messages[0]
367+
summary_message = agent.messages[1]
365368
assert summary_message["role"] == "assistant"
366369

367370
# Extract summary text

0 commit comments

Comments
 (0)