diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py index 216a3e7e15..d95b99a7e7 100644 --- a/astrbot/core/agent/context/manager.py +++ b/astrbot/core/agent/context/manager.py @@ -1,4 +1,5 @@ from astrbot import logger +from astrbot.core.platform.astr_message_event import AstrMessageEvent from ..message import Message from .compressor import LLMSummaryCompressor, TruncateByTurnsCompressor @@ -42,12 +43,17 @@ def __init__( ) async def process( - self, messages: list[Message], trusted_token_usage: int = 0 + self, + messages: list[Message], + trusted_token_usage: int = 0, + event: AstrMessageEvent | None = None, ) -> list[Message]: """Process the messages. Args: messages: The original message list. + trusted_token_usage: Trusted token usage from conversation. + event: Optional event for triggering hooks. Returns: The processed message list. @@ -72,7 +78,7 @@ async def process( if self.compressor.should_compress( result, total_tokens, self.config.max_context_tokens ): - result = await self._run_compression(result, total_tokens) + result = await self._run_compression(result, total_tokens, event) return result except Exception as e: @@ -80,7 +86,10 @@ async def process( return messages async def _run_compression( - self, messages: list[Message], prev_tokens: int + self, + messages: list[Message], + prev_tokens: int, + event: AstrMessageEvent | None = None, ) -> list[Message]: """ Compress/truncate the messages. @@ -88,12 +97,28 @@ async def _run_compression( Args: messages: The original message list. prev_tokens: The token count before compression. + event: Optional event for triggering hooks. Returns: The compressed/truncated message list. """ logger.debug("Compress triggered, starting compression...") + # Trigger before compression hook + if event: + try: + from astrbot.core.pipeline.context_utils import call_event_hook + from astrbot.core.star.star_handler import EventType + + await call_event_hook( + event, + EventType.OnBeforeContextCompressionEvent, + messages, + prev_tokens, + ) + except Exception as e: + logger.warning(f"Hook OnBeforeContextCompressionEvent failed: {e}") + messages = await self.compressor(messages) # double check @@ -117,4 +142,22 @@ async def _run_compression( # still need compress, truncate by half messages = self.truncator.truncate_by_halving(messages) + # Recalculate token count after all truncation steps + final_tokens = self.token_counter.count_tokens(messages) + + # Trigger after compression hook + if event: + try: + from astrbot.core.pipeline.context_utils import call_event_hook + from astrbot.core.star.star_handler import EventType + + await call_event_hook( + event, + EventType.OnAfterContextCompressionEvent, + messages, + final_tokens, + ) + except Exception as e: + logger.warning(f"Hook OnAfterContextCompressionEvent failed: {e}") + return messages diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 9f80dae1c9..1062e87a1e 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -296,8 +296,9 @@ async def step(self): # do truncate and compress token_usage = self.req.conversation.token_usage if self.req.conversation else 0 self._simple_print_message_role("[BefCompact]") + event = getattr(self.run_context.context, "event", None) self.run_context.messages = await self.context_manager.process( - self.run_context.messages, trusted_token_usage=token_usage + self.run_context.messages, trusted_token_usage=token_usage, event=event ) self._simple_print_message_role("[AftCompact]") diff --git a/astrbot/core/star/star_handler.py b/astrbot/core/star/star_handler.py index 63b0c447de..75ea448125 100644 --- a/astrbot/core/star/star_handler.py +++ b/astrbot/core/star/star_handler.py @@ -201,6 +201,8 @@ class EventType(enum.Enum): OnLLMToolRespondEvent = enum.auto() # 调用函数工具后 OnAfterMessageSentEvent = enum.auto() # 发送消息后 OnPluginErrorEvent = enum.auto() # 插件处理消息异常时 + OnBeforeContextCompressionEvent = enum.auto() # 上下文压缩前事件 + OnAfterContextCompressionEvent = enum.auto() # 上下文压缩后事件 H = TypeVar("H", bound=Callable[..., Any])