diff --git a/reme_ai/context/offload/context_offload_op.py b/reme_ai/context/offload/context_offload_op.py index 62086ffc..1c086249 100644 --- a/reme_ai/context/offload/context_offload_op.py +++ b/reme_ai/context/offload/context_offload_op.py @@ -16,6 +16,8 @@ from flowllm.core.op import BaseAsyncOp from flowllm.core.schema import Message from loguru import logger +from reme_ai.enumeration import ContextManageEnum + @C.register_op() class ContextOffloadOp(BaseAsyncOp): @@ -28,45 +30,67 @@ class ContextOffloadOp(BaseAsyncOp): applies additional LLM-based compression to further reduce token count. Context Parameters: + context_manage_mode (ContextManageEnum): The context management mode to use. + - COMPACT: Only applies context compaction to tool messages. + - COMPRESS: Only applies LLM-based compression to messages. + - AUTO: Applies compaction first, then compression if compaction ratio exceeds threshold. + Defaults to AUTO. compact_ratio_threshold (float): Threshold for compaction ratio above which compression - is applied. Defaults to 0.75. If the ratio of compressed tokens to original tokens - exceeds this value, compression will be triggered. + is applied. Only used in AUTO mode. Defaults to 0.75. If the ratio of compressed + tokens to original tokens exceeds this value, compression will be triggered. """ async def async_execute(self): """ Execute the context offload operation. - The operation performs the following steps: - 1. Applies context compaction to reduce token usage in tool messages - 2. Calculates the compaction ratio (compressed tokens / original tokens) - 3. If the ratio exceeds the threshold, applies LLM-based compression to further - reduce token count + The operation behavior depends on the context_manage_mode: + - COMPACT: Only applies context compaction to reduce token usage in tool messages. + - COMPRESS: Only applies LLM-based compression to generate concise summaries. + - AUTO: Applies compaction first, then compression if compaction ratio exceeds threshold. The compaction operation stores full tool message content in external files and - keeps only previews in the context. If this doesn't sufficiently reduce tokens, - the compression operation uses LLM to generate concise summaries of older messages. + keeps only previews in the context. The compression operation uses LLM to generate + concise summaries of older messages. """ from .context_compact_op import ContextCompactOp from .context_compress_op import ContextCompressOp + # Get the context management mode from context, default to AUTO + context_manage_mode = self.context.get("context_manage_mode", ContextManageEnum.AUTO) + if isinstance(context_manage_mode, str): + context_manage_mode = ContextManageEnum(context_manage_mode) + context_compact_op = ContextCompactOp() context_compress_op = ContextCompressOp() - await context_compact_op.async_call(context=self.context) - - origin_messages = [Message(**x) for x in self.context.messages] - origin_token_cnt = self.token_count(origin_messages) - - result_messages = [Message(**x) for x in self.context.response.answer] - answer_token_cnt = self.token_count(result_messages) - - compact_ratio = answer_token_cnt / origin_token_cnt - - compact_ratio_threshold: float = self.context.get("compact_ratio_threshold", 0.75) - if compact_ratio > compact_ratio_threshold: - logger.info(f"Compact ratio {compact_ratio:.2f} > {compact_ratio_threshold:.2f}, compress answer") + if context_manage_mode == ContextManageEnum.COMPACT: + # Only apply compaction + logger.info("Context management mode: COMPACT") + await context_compact_op.async_call(context=self.context) + elif context_manage_mode == ContextManageEnum.COMPRESS: + # Only apply compression + logger.info("Context management mode: COMPRESS") await context_compress_op.async_call(context=self.context) + elif context_manage_mode == ContextManageEnum.AUTO: + # Apply compaction first, then compression if needed + logger.info("Context management mode: AUTO") + await context_compact_op.async_call(context=self.context) + + origin_messages = [Message(**x) for x in self.context.messages] + origin_token_cnt = self.token_count(origin_messages) + + result_messages = [Message(**x) for x in self.context.response.answer] + answer_token_cnt = self.token_count(result_messages) + + compact_ratio = answer_token_cnt / origin_token_cnt + + compact_ratio_threshold: float = self.context.get("compact_ratio_threshold", 0.75) + if compact_ratio > compact_ratio_threshold: + logger.info(f"Compact ratio {compact_ratio:.2f} > {compact_ratio_threshold:.2f}, compress answer") + await context_compress_op.async_call(context=self.context) + else: + raise ValueError(f"Unknown context management mode: {context_manage_mode}") async def async_default_execute(self, e: Exception = None, **_kwargs): """Handle execution errors by returning original messages. diff --git a/reme_ai/enumeration/__init__.py b/reme_ai/enumeration/__init__.py index a6466cf3..482d29e1 100644 --- a/reme_ai/enumeration/__init__.py +++ b/reme_ai/enumeration/__init__.py @@ -4,8 +4,10 @@ This module provides enumerations used throughout the ReMe system, including language enumerations and other type definitions. """ +from reme_ai.enumeration.context_manage_enum import ContextManageEnum from reme_ai.enumeration.language_enum import LanguageEnum __all__ = [ + "ContextManageEnum", "LanguageEnum", ] diff --git a/reme_ai/enumeration/context_manage_enum.py b/reme_ai/enumeration/context_manage_enum.py new file mode 100644 index 00000000..3c290369 --- /dev/null +++ b/reme_ai/enumeration/context_manage_enum.py @@ -0,0 +1,21 @@ +"""Context management enumeration module. + +This module provides enumerations for context management strategies in the ReMe system. +""" + +from enum import Enum + + +class ContextManageEnum(str, Enum): + """ + An enumeration representing context management strategies. + + Members: + - COMPACT: Represents the compact context management strategy. + - COMPRESS: Represents the compress context management strategy. + - AUTO: Represents the automatic context management strategy. + """ + + COMPACT = "compact" + COMPRESS = "compress" + AUTO = "auto"