8000 feat: Update SlidingWindowConversationManager by Unshure · Pull Request #120 · strands-agents/sdk-python · GitHub
[go: up one dir, main page]
Skip to content

feat: Update SlidingWindowConversationManager #120

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

Merged
merged 1 commit into from
May 26, 2025
Merged

feat: Update SlidingWindowConversationManager #120

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
71 changes: 21 additions & 50 deletions src/strands/agent/conversation_manager/sliding_window_conversation_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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__)
Expand Down Expand Up @@ -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.
Expand All @@ -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
2 changes: 1 addition & 1 deletion tests/strands/agent/test_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
55 changes: 17 additions & 38 deletions tests/strands/agent/test_conversation_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
[
Expand All @@ -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},
[
Expand All @@ -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},
[
Expand All @@ -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"}]},
Expand All @@ -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()
Expand Down
0