杜宇轩
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
import json
|
|
from typing import Any
|
|
|
|
from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager
|
|
|
|
DEFAULT_MAX_BYTES = 300000
|
|
DEFAULT_MAX_WINDOW = 400
|
|
|
|
|
|
class ByteWindowConversationManager(SlidingWindowConversationManager):
|
|
def __init__(
|
|
self,
|
|
max_bytes: int = DEFAULT_MAX_BYTES,
|
|
max_window: int = DEFAULT_MAX_WINDOW,
|
|
**kwargs: Any
|
|
):
|
|
super().__init__(max_window=max_window, **kwargs)
|
|
self.max_bytes = max_bytes
|
|
|
|
def _byte_window(self, messages: list) -> int:
|
|
"""计算保留的消息数量"""
|
|
total = 0
|
|
keep = 0
|
|
for m in reversed(messages):
|
|
total += len(json.dumps(m, ensure_ascii=False, default=str).encode("utf-8"))
|
|
if keep >= 1 and total > self.max_bytes:
|
|
break
|
|
keep += 1
|
|
return max(keep, 2)
|
|
|
|
def apply_management(self, agent: "Any", **kwargs: Any) -> None:
|
|
"""应用消息管理策略，确保消息总字节数不超过 max_bytes。"""
|
|
messages = agent.messages
|
|
eff = min(self.window_size, self._byte_window(messages))
|
|
|
|
if len(messages) <= eff:
|
|
return
|
|
|
|
saved = self.window_size
|
|
self.window_size = eff
|
|
try:
|
|
self.reduce_context(agent)
|
|
finally:
|
|
self.window_size = saved
|