From a1f8fbb0f6638d988ecae9d53bbcb5e83767a498 Mon Sep 17 00:00:00 2001 From: Nick Clegg Date: Mon, 26 May 2025 16:33:32 +0000 Subject: [PATCH] feat: Update SlidingWindowConversationManager --- .../sliding_window_conversation_manager.py | 71 ++++++------------- tests/strands/agent/test_agent.py | 2 +- .../agent/test_conversation_manager.py | 55 +++++--------- 3 files changed, 39 insertions(+), 89 deletions(-) diff --git a/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py b/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py index 4b11e81c..f367b272 100644 --- a/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py +++ b/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py @@ -1,12 +1,10 @@ """Sliding window conversation history management.""" -import json import logging -from typing import List, Optional, cast +from typing import Optional -from ...types.content import ContentBlock, Message, Messages +from ...types.content import Message, Messages from ...types.exceptions import ContextWindowOverflowException -from ...types.tools import ToolResult from .conversation_manager import ConversationManager logger = logging.getLogger(__name__) @@ -110,8 +108,9 @@ def _remove_dangling_messages(self, messages: Messages) -> None: def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> None: """Trim the oldest messages to reduce the conversation context size. - The method handles special cases where tool results need to be converted to regular content blocks to maintain - conversation coherence after trimming. + The method handles special cases where trimming the messages leads to: + - toolResult with no corresponding toolUse + - toolUse with no corresponding toolResult Args: messages: The messages to reduce. @@ -126,52 +125,24 @@ def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> N # If the number of messages is less than the window_size, then we default to 2, otherwise, trim to window size trim_index = 2 if len(messages) <= self.window_size else len(messages) - self.window_size - # Throw if we cannot trim any messages from the conversation - if trim_index >= len(messages): - raise ContextWindowOverflowException("Unable to trim conversation context!") from e - - # If the message at the cut index has ToolResultContent, then we map that to ContentBlock. This gets around the - # limitation of needing ToolUse and ToolResults to be paired. - if any("toolResult" in content for content in messages[trim_index]["content"]): - if len(messages[trim_index]["content"]) == 1: - messages[trim_index]["content"] = self._map_tool_result_content( - cast(ToolResult, messages[trim_index]["content"][0]["toolResult"]) + # Find the next valid trim_index + while trim_index < len(messages): + if ( + # Oldest message cannot be a toolResult because it needs a toolUse preceding it + any("toolResult" in content for content in messages[trim_index]["content"]) + or ( + # Oldest message can be a toolUse only if a toolResult immediately follows it. + any("toolUse" in content for content in messages[trim_index]["content"]) + and trim_index + 1 < len(messages) + and not any("toolResult" in content for content in messages[trim_index + 1]["content"]) ) - - # If there is more content than just one ToolResultContent, then we cannot cut at this index. + ): + trim_index += 1 else: - raise ContextWindowOverflowException("Unable to trim conversation context!") from e + break + else: + # If we didn't find a valid trim_index, then we throw + raise ContextWindowOverflowException("Unable to trim conversation context!") from e # Overwrite message history messages[:] = messages[trim_index:] - - def _map_tool_result_content(self, tool_result: ToolResult) -> List[ContentBlock]: - """Convert a ToolResult to a list of standard ContentBlocks. - - This method transforms tool result content into standard content blocks that can be preserved when trimming the - conversation history. - - Args: - tool_result: The ToolResult to convert. - - Returns: - A list of content blocks representing the tool result. - """ - contents = [] - text_content = "Tool Result Status: " + tool_result["status"] if tool_result["status"] else "" - - for tool_result_content in tool_result["content"]: - if "text" in tool_result_content: - text_content = "\nTool Result Text Content: " + tool_result_content["text"] + f"\n{text_content}" - elif "json" in tool_result_content: - text_content = ( - "\nTool Result JSON Content: " + json.dumps(tool_result_content["json"]) + f"\n{text_content}" - ) - elif "image" in tool_result_content: - contents.append(ContentBlock(image=tool_result_content["image"])) - elif "document" in tool_result_content: - contents.append(ContentBlock(document=tool_result_content["document"])) - else: - logger.warning("unsupported content type") - contents.append(ContentBlock(text=text_content)) - return contents diff --git a/tests/strands/agent/test_agent.py b/tests/strands/agent/test_agent.py index 5c7d11e4..85f512b0 100644 --- a/tests/strands/agent/test_agent.py +++ b/tests/strands/agent/test_agent.py @@ -428,7 +428,7 @@ def test_agent__call__always_sliding_window_conversation_manager_doesnt_infinite with pytest.raises(ContextWindowOverflowException): agent("Test!") - assert conversation_manager_spy.reduce_context.call_count == 251 + assert conversation_manager_spy.reduce_context.call_count > 0 assert conversation_manager_spy.apply_management.call_count == 1 diff --git a/tests/strands/agent/test_conversation_manager.py b/tests/strands/agent/test_conversation_manager.py index 2f6ee77d..b6132f1d 100644 --- a/tests/strands/agent/test_conversation_manager.py +++ b/tests/strands/agent/test_conversation_manager.py @@ -111,41 +111,7 @@ def conversation_manager(request): {"role": "assistant", "content": [{"text": "Second response"}]}, ], ), - # 7 - Message count above max window size - Remove dangling tool uses and tool results - ( - {"window_size": 1}, - [ - {"role": "user", "content": [{"text": "First message"}]}, - {"role": "assistant", "content": [{"toolUse": {"toolUseId": "321", "name": "tool1", "input": {}}}]}, - { - "role": "user", - "content": [ - {"toolResult": {"toolUseId": "123", "content": [{"text": "Hello!"}], "status": "success"}} - ], - }, - ], - [ - { - "role": "user", - "content": [{"text": "\nTool Result Text Content: Hello!\nTool Result Status: success"}], - }, - ], - ), - # 8 - Message count above max window size - Remove multiple tool use/tool result pairs - ( - {"window_size": 1}, - [ - {"role": "user", "content": [{"toolResult": {"toolUseId": "123", "content": [], "status": "success"}}]}, - {"role": "assistant", "content": [{"toolUse": {"toolUseId": "123", "name": "tool1", "input": {}}}]}, - {"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]}, - {"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool1", "input": {}}}]}, - {"role": "user", "content": [{"toolResult": {"toolUseId": "789", "content": [], "status": "success"}}]}, - ], - [ - {"role": "user", "content": [{"text": "Tool Result Status: success"}]}, - ], - ), - # 9 - Message count above max window size - Preserve tool use/tool result pairs + # 7 - Message count above max window size - Preserve tool use/tool result pairs ( {"window_size": 2}, [ @@ -158,7 +124,7 @@ def conversation_manager(request): {"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]}, ], ), - # 10 - Test sliding window behavior - preserve tool use/result pairs across cut boundary + # 8 - Test sliding window behavior - preserve tool use/result pairs across cut boundary ( {"window_size": 3}, [ @@ -173,7 +139,7 @@ def conversation_manager(request): {"role": "assistant", "content": [{"text": "Response after tool use"}]}, ], ), - # 11 - Test sliding window with multiple tool pairs that need preservation + # 9 - Test sliding window with multiple tool pairs that need preservation ( {"window_size": 4}, [ @@ -185,7 +151,6 @@ def conversation_manager(request): {"role": "assistant", "content": [{"text": "Final response"}]}, ], [ - {"role": "user", "content": [{"text": "Tool Result Status: success"}]}, {"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool2", "input": {}}}]}, {"role": "user", "content": [{"toolResult": {"toolUseId": "456", "content": [], "status": "success"}}]}, {"role": "assistant", "content": [{"text": "Final response"}]}, @@ -200,6 +165,20 @@ def test_apply_management(conversation_manager, messages, expected_messages): assert messages == expected_messages +def test_sliding_window_conversation_manager_with_untrimmable_history_raises_context_window_overflow_exception(): + manager = strands.agent.conversation_manager.SlidingWindowConversationManager(1) + messages = [ + {"role": "assistant", "content": [{"toolUse": {"toolUseId": "456", "name": "tool1", "input": {}}}]}, + {"role": "user", "content": [{"toolResult": {"toolUseId": "789", "content": [], "status": "success"}}]}, + ] + original_messages = messages.copy() + + with pytest.raises(ContextWindowOverflowException): + manager.apply_management(messages) + + assert messages == original_messages + + def test_null_conversation_manager_reduce_context_raises_context_window_overflow_exception(): """Test that NullConversationManager doesn't modify messages.""" manager = strands.agent.conversation_manager.NullConversationManager()