feat(summarizing_conversation_manager): implement summarization of older messages

stefanoamorelli · stefanoamorelli · commit bca4e6e9ad0d · 2025-05-27T21:33:21.000+03:00
diff --git a/src/strands/agent/conversation_manager/__init__.py b/src/strands/agent/conversation_manager/__init__.py
@@ -6,6 +6,8 @@
 - NullConversationManager: A no-op implementation that does not modify conversation history
 - SlidingWindowConversationManager: An implementation that maintains a sliding window of messages to control context
   size while preserving conversation coherence
+- SummarizingConversationManager: An extension of sliding window that can optionally summarize older context instead
+  of simply trimming it
 
 Conversation managers help control memory usage and context length while maintaining relevant conversation state, which
 is critical for effective agent interactions.
@@ -14,5 +16,11 @@
 from .conversation_manager import ConversationManager
 from .null_conversation_manager import NullConversationManager
 from .sliding_window_conversation_manager import SlidingWindowConversationManager
+from .summarizing_conversation_manager import SummarizingConversationManager
 
-__all__ = ["ConversationManager", "NullConversationManager", "SlidingWindowConversationManager"]
+__all__ = [
+    "ConversationManager",
+    "NullConversationManager",
+    "SlidingWindowConversationManager",
+    "SummarizingConversationManager",
+]
diff --git a/src/strands/agent/conversation_manager/conversation_manager.py b/src/strands/agent/conversation_manager/conversation_manager.py
@@ -1,10 +1,13 @@
 """Abstract interface for conversation history management."""
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from ...types.content import Messages
 
+if TYPE_CHECKING:
+    from ..agent import Agent
+
 
 class ConversationManager(ABC):
     """Abstract base class for managing conversation history.
@@ -34,7 +37,9 @@ def apply_management(self, messages: Messages) -> None:
 
     @abstractmethod
     # pragma: no cover
-    def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> None:
+    def reduce_context(
+        self, messages: Messages, e: Optional[Exception] = None, agent: Optional["Agent"] = None
+    ) -> None:
         """Called when the model's context window is exceeded.
 
         This method should implement the specific strategy for reducing the window size when a context overflow occurs.
@@ -51,5 +56,6 @@ def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> N
             messages: The conversation history to reduce.
                 This list is modified in-place.
             e: The exception that triggered the context reduction, if any.
+            agent: The agent instance, if available for advanced reduction strategies.
         """
         pass
diff --git a/src/strands/agent/conversation_manager/null_conversation_manager.py b/src/strands/agent/conversation_manager/null_conversation_manager.py
@@ -25,12 +25,13 @@ def apply_management(self, messages: Messages) -> None:
         """
         pass
 
-    def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> None:
+    def reduce_context(self, messages: Messages, e: Optional[Exception] = None, **kwargs) -> None:
         """Does not reduce context and raises an exception.
 
         Args:
             messages: The conversation history that will remain unmodified.
             e: The exception that triggered the context reduction, if any.
+            **kwargs: Additional keyword arguments (ignored).
 
         Raises:
             e: If provided.
diff --git a/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py b/src/strands/agent/conversation_manager/sliding_window_conversation_manager.py
@@ -105,7 +105,7 @@ def _remove_dangling_messages(self, messages: Messages) -> None:
                     if not any("toolResult" in content for content in messages[-1]["content"]):
                         messages.pop()
 
-    def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> None:
+    def reduce_context(self, messages: Messages, e: Optional[Exception] = None, **kwargs) -> None:
         """Trim the oldest messages to reduce the conversation context size.
 
         The method handles special cases where trimming the messages leads to:
@@ -116,6 +116,7 @@ def reduce_context(self, messages: Messages, e: Optional[Exception] = None) -> N
             messages: The messages to reduce.
                 This list is modified in-place.
             e: The exception that triggered the context reduction, if any.
+            **kwargs: Additional keyword arguments (ignored).
 
         Raises:
             ContextWindowOverflowException: If the context cannot be reduced further.
diff --git a/src/strands/agent/conversation_manager/summarizing_conversation_manager.py b/src/strands/agent/conversation_manager/summarizing_conversation_manager.py
@@ -0,0 +1,232 @@
+"""Summarizing conversation history management with configurable options."""
+
+import json
+import logging
+from typing import TYPE_CHECKING, List, Optional, cast
+
+from ...types.content import ContentBlock, Message, Messages
+from ...types.exceptions import ContextWindowOverflowException
+from ...types.tools import ToolResult
+from .conversation_manager import ConversationManager
+
+if TYPE_CHECKING:
+    from ..agent import Agent
+
+
+logger = logging.getLogger(__name__)
+
+
+class SummarizingConversationManager(ConversationManager):
+    """Extends sliding window manager with optional summarization of older messages.
+
+    This manager provides a configurable option to summarize older context instead of
+    simply trimming it, helping preserve important information while staying within
+    context limits.
+    """
+
+    def __init__(
+        self,
+        window_size: int = 40,
+        summary_ratio: float = 0.3,
+        preserve_recent_messages: int = 10,
+        summarization_agent: Optional["Agent"] = None,
+        summarization_system_prompt: Optional[str] = None,
+    ):
+        """Initialize the summarizing conversation manager.
+
+        Args:
+            window_size: Maximum number of messages to keep in history.
+                Defaults to 40 messages.
+            summary_ratio: Ratio of messages to summarize vs keep when window is exceeded.
+                Value between 0.1 and 0.8. Defaults to 0.3 (summarize 30% of oldest messages).
+            preserve_recent_messages: Minimum number of recent messages to always keep.
+                Defaults to 10 messages.
+            summarization_agent: Optional agent to use for summarization instead of the parent agent.
+                If provided, this agent can use tools as part of the summarization process.
+            summarization_system_prompt: Optional system prompt override for summarization.
+                If None, uses the default summarization prompt.
+        """
+        self.window_size = window_size
+        self.summary_ratio = max(0.1, min(0.8, summary_ratio))
+        self.preserve_recent_messages = preserve_recent_messages
+        self.summarization_agent = summarization_agent
+        self.summarization_system_prompt = summarization_system_prompt
+
+    def apply_management(self, messages: Messages) -> None:
+        """Apply management strategy when message count exceeds window size.
+
+        Args:
+            messages: The conversation history to manage.
+                This list is modified in-place.
+        """
+        if len(messages) > self.window_size:
+            # Call reduce_context (agent will be provided by the caller when needed)
+            self.reduce_context(messages)
+
+    def reduce_context(
+        self, messages: Messages, e: Optional[Exception] = None, agent: Optional["Agent"] = None
+    ) -> None:
+        """Reduce context using summarization.
+
+        Args:
+            messages: The messages to reduce.
+                This list is modified in-place.
+            e: The exception that triggered the context reduction, if any.
+            agent: The agent instance to use for summarization.
+                If None, context overflow will raise an exception.
+
+        Raises:
+            ContextWindowOverflowException: If the context cannot be summarized.
+        """
+        if agent is None:
+            raise ContextWindowOverflowException("No agent provided for context reduction") from e
+
+        try:
+            # Calculate how many messages to summarize
+            messages_to_summarize_count = max(1, int(len(messages) * self.summary_ratio))
+
+            # Ensure we don't summarize recent messages
+            messages_to_summarize_count = min(
+                messages_to_summarize_count, len(messages) - self.preserve_recent_messages
+            )
+
+            if messages_to_summarize_count <= 0:
+                raise ContextWindowOverflowException("Cannot summarize: insufficient messages for summarization") from e
+
+            # Extract messages to summarize
+            messages_to_summarize = messages[:messages_to_summarize_count]
+            remaining_messages = messages[messages_to_summarize_count:]
+
+            # Generate summary
+            summary_content = self._generate_summary(messages_to_summarize, agent)
+
+            # Create a summary message
+            summary_message: Message = {"role": "assistant", "content": [ContentBlock(text=summary_content)]}
+
+            # Replace the summarized messages with the summary
+            messages[:] = [summary_message] + remaining_messages
+
+        except Exception as summarization_error:
+            logger.error("Summarization failed: %s", summarization_error)
+            raise
+
+    def _generate_summary(self, messages: List[Message], agent: "Agent") -> str:
+        """Generate a summary of the provided messages.
+
+        Args:
+            messages: The messages to summarize.
+            agent: The agent instance to use for summarization.
+
+        Returns:
+            A text summary of the conversation history.
+
+        Raises:
+            Exception: If summary generation fails.
+        """
+        # Choose which agent to use for summarization
+        summarization_agent = self.summarization_agent if self.summarization_agent is not None else agent
+
+        # Format messages for summarization, preserving rich content
+        formatted_messages = self._format_messages_for_summarization(messages)
+
+        # Use custom system prompt if provided, otherwise use default
+        system_prompt = (
+            self.summarization_system_prompt
+            if self.summarization_system_prompt is not None
+            else (
+                "You are a conversation summarizer. Provide a concise summary of the conversation history.\n\n"
+                "Format Requirements:\n"
+                "- You MUST create a structured and concise summary in bullet-point format.\n"
+                "- You MUST NOT respond conversationally.\n"
+                "- You MUST NOT address the user directly.\n\n"
+                "Task:\n"
+                "Your task is to create a structured summary document:\n"
+                "- It MUST contain bullet points with key topics and questions covered\n"
+                "- It MUST contain bullet points for all significant tools executed and their results\n"
+                "- It MUST contain bullet points for any code or technical information shared\n"
+                "- It MUST contain a section of key insights gained\n"
+                "- It MUST format the summary in the third person\n\n"
+                "Example format:\n\n"
+                "## Conversation Summary\n"
+                "* Topic 1: Key information\n"
+                "* Topic 2: Key information\n"
+                "*\n"
+                "## Tools Executed\n"
+                "* Tool X: Result Y"
+            )
+        )
+
+        # Save original system prompt and messages to restore later
+        original_system_prompt = summarization_agent.system_prompt
+        original_messages = summarization_agent.messages.copy()
+
+        try:
+            # Temporarily set the system prompt and set formatted messages for summarization
+            summarization_agent.system_prompt = system_prompt
+            summarization_agent.messages = formatted_messages
+
+            # Use the agent to generate summary with rich content (can use tools if needed)
+            result = summarization_agent("Please summarize this conversation.")
+            summary_response = ""
+
+            # Extract text content from the result
+            if hasattr(result, "message") and result.message and "content" in result.message:
+                for content_block in result.message["content"]:
+                    if "text" in content_block:
+                        summary_response += content_block["text"]
+
+            return summary_response.strip()
+
+        finally:
+            # Restore original agent state
+            summarization_agent.system_prompt = original_system_prompt
+            summarization_agent.messages = original_messages
+
+    def _format_messages_for_summarization(self, messages: List[Message]) -> List[Message]:
+        """Format messages for summarization, preserving rich content.
+
+        Args:
+            messages: The messages to format.
+
+        Returns:
+            A list of formatted messages with rich content preserved for agent analysis.
+        """
+        formatted_messages = []
+
+        for _, message in enumerate(messages):
+            role = message["role"]
+            formatted_content = []
+
+            for content in message["content"]:
+                if "text" in content:
+                    formatted_content.append(ContentBlock(text=content["text"]))
+                elif "toolUse" in content:
+                    tool_use = content["toolUse"]
+                    formatted_content.append(ContentBlock(text=f"[Used tool: {tool_use.get('name', 'unknown')}]"))
+                elif "toolResult" in content:
+                    tool_result = cast(ToolResult, content["toolResult"])
+                    status = tool_result.get("status", "")
+                    if tool_result.get("content"):
+                        for tr_content in tool_result["content"]:
+                            if "text" in tr_content:
+                                text = tr_content["text"]
+                                formatted_content.append(ContentBlock(text=f"[Tool result text: {text}]"))
+                            elif "json" in tr_content:
+                                json_str = json.dumps(tr_content["json"])
+                                formatted_content.append(ContentBlock(text=f"[Tool result JSON: {json_str}]"))
+                            elif "image" in tr_content:
+                                formatted_content.append(ContentBlock(image=tr_content["image"]))
+                            elif "document" in tr_content:
+                                formatted_content.append(ContentBlock(document=tr_content["document"]))
+                    else:
+                        formatted_content.append(ContentBlock(text=f"[Tool result: {status}]"))
+                elif "image" in content:
+                    formatted_content.append(ContentBlock(image=content["image"]))
+                elif "document" in content:
+                    formatted_content.append(ContentBlock(document=content["document"]))
+
+            if formatted_content:
+                formatted_message = {"role": role, "content": formatted_content}
+                formatted_messages.append(formatted_message)
+
+        return formatted_messages
diff --git a/tests/strands/agent/test_summarizing_conversation_manager.py b/tests/strands/agent/test_summarizing_conversation_manager.py
