|
| 1 | +"""Summarizing conversation history management with configurable options.""" |
| 2 | + |
| 3 | +import logging |
| 4 | +from typing import TYPE_CHECKING, List, Optional |
| 5 | + |
| 6 | +from ...types.content import Message, Messages |
| 7 | +from ...types.exceptions import ContextWindowOverflowException |
| 8 | +from .conversation_manager import ConversationManager |
| 9 | + |
| 10 | +if TYPE_CHECKING: |
| 11 | + from ..agent import Agent |
| 12 | + |
| 13 | + |
| 14 | +logger = logging.getLogger(__name__) |
| 15 | + |
| 16 | + |
| 17 | +DEFAULT_SUMMARIZATION_PROMPT = """You are a conversation summarizer. Provide a concise summary of the conversation \ |
| 18 | +history. |
| 19 | +
|
| 20 | +Format Requirements: |
| 21 | +- You MUST create a structured and concise summary in bullet-point format. |
| 22 | +- You MUST NOT respond conversationally. |
| 23 | +- You MUST NOT address the user directly. |
| 24 | +
|
| 25 | +Task: |
| 26 | +Your task is to create a structured summary document: |
| 27 | +- It MUST contain bullet points with key topics and questions covered |
| 28 | +- It MUST contain bullet points for all significant tools executed and their results |
| 29 | +- It MUST contain bullet points for any code or technical information shared |
| 30 | +- It MUST contain a section of key insights gained |
| 31 | +- It MUST format the summary in the third person |
| 32 | +
|
| 33 | +Example format: |
| 34 | +
|
| 35 | +## Conversation Summary |
| 36 | +* Topic 1: Key information |
| 37 | +* Topic 2: Key information |
| 38 | +* |
| 39 | +## Tools Executed |
| 40 | +* Tool X: Result Y""" |
| 41 | + |
| 42 | + |
| 43 | +class SummarizingConversationManager(ConversationManager): |
| 44 | + """Extends sliding window manager with optional summarization of older messages. |
| 45 | +
|
| 46 | + This manager provides a configurable option to summarize older context instead of |
| 47 | + simply trimming it, helping preserve important information while staying within |
| 48 | + context limits. |
| 49 | + """ |
| 50 | + |
| 51 | + def __init__( |
| 52 | + self, |
| 53 | + summary_ratio: float = 0.3, |
| 54 | + preserve_recent_messages: int = 10, |
| 55 | + summarization_agent: Optional["Agent"] = None, |
| 56 | + summarization_system_prompt: Optional[str] = None, |
| 57 | + ): |
| 58 | + """Initialize the summarizing conversation manager. |
| 59 | +
|
| 60 | + Args: |
| 61 | + summary_ratio: Ratio of messages to summarize vs keep when context overflow occurs. |
| 62 | + Value between 0.1 and 0.8. Defaults to 0.3 (summarize 30% of oldest messages). |
| 63 | + preserve_recent_messages: Minimum number of recent messages to always keep. |
| 64 | + Defaults to 10 messages. |
| 65 | + summarization_agent: Optional agent to use for summarization instead of the parent agent. |
| 66 | + If provided, this agent can use tools as part of the summarization process. |
| 67 | + summarization_system_prompt: Optional system prompt override for summarization. |
| 68 | + If None, uses the default summarization prompt. |
| 69 | + """ |
| 70 | + self.summary_ratio = max(0.1, min(0.8, summary_ratio)) |
| 71 | + self.preserve_recent_messages = preserve_recent_messages |
| 72 | + self.summarization_agent = summarization_agent |
| 73 | + self.summarization_system_prompt = summarization_system_prompt |
| 74 | + |
| 75 | + def apply_management(self, messages: Messages) -> None: |
| 76 | + """Apply management strategy to conversation history. |
| 77 | +
|
| 78 | + For the summarizing conversation manager, no proactive management is performed. |
| 79 | + Summarization only occurs when there's a context overflow that triggers reduce_context. |
| 80 | +
|
| 81 | + Args: |
| 82 | + messages: The conversation history to manage. |
| 83 | + This list is modified in-place. |
| 84 | + """ |
| 85 | + # No proactive management - summarization only happens on context overflow |
| 86 | + pass |
| 87 | + |
| 88 | + def reduce_context( |
| 89 | + self, messages: Messages, e: Optional[Exception] = None, agent: Optional["Agent"] = None |
| 90 | + ) -> None: |
| 91 | + """Reduce context using summarization. |
| 92 | +
|
| 93 | + Args: |
| 94 | + messages: The messages to reduce. |
| 95 | + This list is modified in-place. |
| 96 | + e: The exception that triggered the context reduction, if any. |
| 97 | + agent: The agent instance to use for summarization. |
| 98 | + If None, context overflow will raise an exception. |
| 99 | +
|
| 100 | + Raises: |
| 101 | + ContextWindowOverflowException: If the context cannot be summarized. |
| 102 | + """ |
| 103 | + if agent is None: |
| 104 | + raise ContextWindowOverflowException("No agent provided for context reduction") from e |
| 105 | + |
| 106 | + try: |
| 107 | + # Calculate how many messages to summarize |
| 108 | + messages_to_summarize_count = max(1, int(len(messages) * self.summary_ratio)) |
| 109 | + |
| 110 | + # Ensure we don't summarize recent messages |
| 111 | + messages_to_summarize_count = min( |
| 112 | + messages_to_summarize_count, len(messages) - self.preserve_recent_messages |
| 113 | + ) |
| 114 | + |
| 115 | + if messages_to_summarize_count <= 0: |
| 116 | + raise ContextWindowOverflowException("Cannot summarize: insufficient messages for summarization") from e |
| 117 | + |
| 118 | + # Adjust split point to avoid breaking ToolUse/ToolResult pairs |
| 119 | + messages_to_summarize_count = self._adjust_split_point_for_tool_pairs(messages, messages_to_summarize_count) |
| 120 | + |
| 121 | + if messages_to_summarize_count <= 0: |
| 122 | + raise ContextWindowOverflowException("Cannot summarize: insufficient messages for summarization") from e |
| 123 | + |
| 124 | + # Extract messages to summarize |
| 125 | + messages_to_summarize = messages[:messages_to_summarize_count] |
| 126 | + remaining_messages = messages[messages_to_summarize_count:] |
| 127 | + |
| 128 | + # Generate summary |
| 129 | + summary_message = self._generate_summary(messages_to_summarize, agent) |
| 130 | + |
| 131 | + # Replace the summarized messages with the summary |
| 132 | + messages[:] = [summary_message] + remaining_messages |
| 133 | + |
| 134 | + except Exception as summarization_error: |
| 135 | + logger.error("Summarization failed: %s", summarization_error) |
| 136 | + raise |
| 137 | + |
| 138 | + def _generate_summary(self, messages: List[Message], agent: "Agent") -> Message: |
| 139 | + """Generate a summary of the provided messages. |
| 140 | +
|
| 141 | + Args: |
| 142 | + messages: The messages to summarize. |
| 143 | + agent: The agent instance to use for summarization. |
| 144 | +
|
| 145 | + Returns: |
| 146 | + A message containing the conversation summary. |
| 147 | +
|
| 148 | + Raises: |
| 149 | + Exception: If summary generation fails. |
| 150 | + """ |
| 151 | + # Choose which agent to use for summarization |
| 152 | + summarization_agent = self.summarization_agent if self.summarization_agent is not None else agent |
| 153 | + |
| 154 | + # Format messages for summarization, preserving rich content |
| 155 | + formatted_messages = self._format_messages_for_summarization(messages) |
| 156 | + |
| 157 | + # Use custom system prompt if provided, otherwise use default |
| 158 | + system_prompt = ( |
| 159 | + self.summarization_system_prompt |
| 160 | + if self.summarization_system_prompt is not None |
| 161 | + else DEFAULT_SUMMARIZATION_PROMPT |
| 162 | + ) |
| 163 | + |
| 164 | + # Save original system prompt and messages to restore later |
| 165 | + original_system_prompt = summarization_agent.system_prompt |
| 166 | + original_messages = summarization_agent.messages.copy() |
| 167 | + |
| 168 | + try: |
| 169 | + # Temporarily set the system prompt and set formatted messages for summarization |
| 170 | + summarization_agent.system_prompt = system_prompt |
| 171 | + summarization_agent.messages = formatted_messages |
| 172 | + |
| 173 | + # Use the agent to generate summary with rich content (can use tools if needed) |
| 174 | + result = summarization_agent("Please summarize this conversation.") |
| 175 | + |
| 176 | + return result.message |
| 177 | + |
| 178 | + finally: |
| 179 | + # Restore original agent state |
| 180 | + summarization_agent.system_prompt = original_system_prompt |
| 181 | + summarization_agent.messages = original_messages |
| 182 | + |
| 183 | + def _adjust_split_point_for_tool_pairs(self, messages: List[Message], split_point: int) -> int: |
| 184 | + """Adjust the split point to avoid breaking ToolUse/ToolResult pairs. |
| 185 | +
|
| 186 | + Args: |
| 187 | + messages: The full list of messages. |
| 188 | + split_point: The initially calculated split point. |
| 189 | +
|
| 190 | + Returns: |
| 191 | + The adjusted split point that doesn't break ToolUse/ToolResult pairs. |
| 192 | + """ |
| 193 | + if split_point >= len(messages): |
| 194 | + return split_point |
| 195 | + |
| 196 | + # Check if the message at split_point is a ToolResult |
| 197 | + if split_point < len(messages): |
| 198 | + message_at_split = messages[split_point] |
| 199 | + if self._message_contains_tool_result(message_at_split): |
| 200 | + # Find the corresponding ToolUse by looking backwards |
| 201 | + for i in range(split_point - 1, -1, -1): |
| 202 | + if self._message_contains_tool_use(messages[i]): |
| 203 | + # Move split point to before the ToolUse |
| 204 | + return i |
| 205 | + # If no ToolUse found, move split point back by 1 |
| 206 | + return max(0, split_point - 1) |
| 207 | + |
| 208 | + return split_point |
| 209 | + |
| 210 | + def _message_contains_tool_use(self, message: Message) -> bool: |
| 211 | + """Check if a message contains a ToolUse.""" |
| 212 | + for content in message.get("content", []): |
| 213 | + if "toolUse" in content: |
| 214 | + return True |
| 215 | + return False |
| 216 | + |
| 217 | + def _message_contains_tool_result(self, message: Message) -> bool: |
| 218 | + """Check if a message contains a ToolResult.""" |
| 219 | + for content in message.get("content", []): |
| 220 | + if "toolResult" in content: |
| 221 | + return True |
| 222 | + return False |
| 223 | + |
| 224 | + def _format_messages_for_summarization(self, messages: List[Message]) -> List[Message]: |
| 225 | + """Format messages for summarization, preserving original content. |
| 226 | +
|
| 227 | + Since Bedrock validation restrictions have been relaxed, we can now pass |
| 228 | + messages with their original content (including ToolUse/ToolResult) directly |
| 229 | + to the summarization agent without lossy conversions. |
| 230 | +
|
| 231 | + Args: |
| 232 | + messages: The messages to format. |
| 233 | +
|
| 234 | + Returns: |
| 235 | + The original messages, preserving all content types and structure. |
| 236 | + """ |
| 237 | + return messages |
0 commit comments