mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-05 02:41:43 +00:00
98 lines
3.9 KiB
Python
98 lines
3.9 KiB
Python
"""ContextChecker module for checking context size and splitting messages."""
|
|
|
|
from agentscope.message import Msg
|
|
from agentscope.token import HuggingFaceTokenCounter
|
|
|
|
from ..utils import AsMsgHandler
|
|
from ....core.op import BaseOp
|
|
from ....core.utils import get_std_logger
|
|
|
|
logger = get_std_logger()
|
|
|
|
|
|
class ContextChecker(BaseOp):
|
|
"""
|
|
ContextChecker class for checking context size and splitting messages.
|
|
|
|
This class analyzes conversation messages to determine if the context
|
|
exceeds the specified token threshold and splits messages into two groups:
|
|
those that should be compacted and those to keep in context.
|
|
|
|
Attributes:
|
|
memory_compact_threshold (int): Token count threshold for triggering compaction.
|
|
memory_compact_reserve (int): Token count to reserve for recent messages.
|
|
msg_handler (AsMsgHandler): Handler for message processing and token counting.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
memory_compact_threshold: int,
|
|
memory_compact_reserve: int = 10000,
|
|
token_counter: HuggingFaceTokenCounter | None = None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Initialize the ContextChecker.
|
|
|
|
Args:
|
|
memory_compact_threshold (int): Token count threshold for triggering
|
|
compaction. Messages exceeding this threshold will be split.
|
|
memory_compact_reserve (int): Token count to reserve for recent messages
|
|
to keep in context. Defaults to 10000 tokens.
|
|
token_counter (HuggingFaceTokenCounter | None): Token counter for
|
|
measuring content length. If None, a default counter will be used.
|
|
**kwargs: Additional keyword arguments passed to BaseOp.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self.memory_compact_threshold: int = memory_compact_threshold
|
|
self.memory_compact_reserve: int = memory_compact_reserve
|
|
assert self.memory_compact_threshold > self.memory_compact_reserve
|
|
|
|
self.msg_handler = AsMsgHandler(token_counter=token_counter)
|
|
|
|
async def execute(self) -> tuple[list[Msg], list[Msg], bool]:
|
|
"""
|
|
Execute context check and split messages.
|
|
|
|
Retrieves messages from context and checks if they exceed the token
|
|
threshold. If so, splits them into messages to compact and messages
|
|
to keep.
|
|
|
|
Context Parameters:
|
|
messages (list[Msg]): List of conversation messages to check.
|
|
Retrieved from self.context.get("messages", []).
|
|
|
|
Returns:
|
|
tuple[list[Msg], list[Msg], bool]: A tuple containing:
|
|
- messages_to_compact (list[Msg]): Older messages that should
|
|
be compacted/summarized.
|
|
- messages_to_keep (list[Msg]): Recent messages to keep in context.
|
|
- is_valid (bool): True if the split is valid (tool calls aligned),
|
|
False if splitting would break conversation integrity.
|
|
|
|
Note:
|
|
- Returns ([], messages, True) if no compaction is needed.
|
|
- Ensures conversation pairs (user-assistant) are not split.
|
|
- is_valid=False indicates tool_use and tool_result are misaligned.
|
|
"""
|
|
messages: list[Msg] = self.context.get("messages", [])
|
|
|
|
if not messages:
|
|
logger.info("ContextChecker: No messages to check.")
|
|
return [], [], True
|
|
|
|
messages_to_compact, messages_to_keep, is_valid = self.msg_handler.context_check(
|
|
messages=messages,
|
|
memory_compact_threshold=self.memory_compact_threshold,
|
|
memory_compact_reserve=self.memory_compact_reserve,
|
|
)
|
|
|
|
if messages_to_compact:
|
|
logger.info(
|
|
f"ContextChecker Result: "
|
|
f"to_compact={len(messages_to_compact)}, "
|
|
f"to_keep={len(messages_to_keep)}, "
|
|
f"is_valid={is_valid}",
|
|
)
|
|
|
|
return messages_to_compact, messages_to_keep, is_valid
|