ReMe/reme_ai/summary/working/message_compact_op.py

"""Working-memory compaction module for reducing token usage.

This module provides functionality to compact large tool messages for
*working memory summary* by storing their full content in external files and
keeping only short previews in the active message list. This reduces token
usage while preserving access to detailed information when needed.
"""

from pathlib import Path
from uuid import uuid4

from flowllm.core.context import C
from flowllm.core.enumeration import Role
from flowllm.core.op import BaseAsyncOp
from flowllm.core.schema import Message
from loguru import logger


@C.register_op()
class MessageCompactOp(BaseAsyncOp):
    """
    Working-memory compaction operation that reduces token usage by compacting tool messages.

    When the total token count exceeds the threshold, this operation truncates large tool
    messages and stores the full content in external files. This is intended for
    working-memory style summarization: the agent sees a short preview while the
    complete result remains available out-of-band.
    """

    async def async_execute(self):
        """
        Execute the context compaction operation.

        The operation:
        1. Calculates the total token count of all messages
        2. If below threshold, returns messages unchanged
        3. Otherwise, compresses large tool messages by:
           - Keeping only a preview of the content
           - Storing full content in external files
           - Preserving recent tool messages
        """
        # Get configuration from context
        max_total_tokens: int = self.context.get("max_total_tokens", 20000)
        max_tool_message_tokens: int = self.context.get("max_tool_message_tokens", 2000)
        preview_char_length: int = self.context.get("preview_char_length", 0)
        keep_recent_count: int = self.context.get("keep_recent_count", 10)
        store_dir: Path = Path(self.context.get("store_dir", ""))

        assert max_total_tokens > 0, "max_total_tokens must be greater than 0"
        assert max_tool_message_tokens > 0, "max_tool_message_tokens must be greater than 0"
        assert preview_char_length >= 0, "preview_char_length must be greater than 0"
        assert keep_recent_count >= 0, "keep_recent_count must be greater than 0"

        messages = [Message(**x) for x in self.context.messages]

        # Convert context messages to Message objects
        messages_to_compress = [x for x in messages if x.role is not Role.SYSTEM]
        if keep_recent_count > 0:
            messages_to_compress = messages_to_compress[:-keep_recent_count]

        # If nothing to compress after filtering, return original messages
        if not messages_to_compress:
            self.context.response.answer = self.context.messages
            logger.info("No messages to compress after filtering, returning original messages")
            return

        logger.info(f"{len(messages_to_compress)} messages remaining for compression check")

        # Calculate total token count
        compact_token_cnt: int = self.token_count(messages_to_compress)
        logger.info(f"Context compaction check: total token count={compact_token_cnt}, threshold={max_total_tokens}")

        # If token count is within threshold, no compaction needed
        if compact_token_cnt <= max_total_tokens:
            self.context.response.answer = self.context.messages
            logger.info(f"Token count ({compact_token_cnt}) is within ({max_total_tokens}), no compaction needed")
            return

        # Filter tool messages for processing
        tool_messages = [x for x in messages_to_compress if x.role is Role.TOOL]

        # Dictionary to store file paths and their full content (for potential batch writing)
        write_file_dict = {}

        # Process each tool message
        for tool_message in tool_messages:
            # Calculate token count for this specific tool message
            tool_token_cnt = self.token_count([tool_message])

            # Skip if token count is within threshold
            if tool_token_cnt <= max_tool_message_tokens:
                logger.info(
                    f"Skipping tool message (tool_call_id={tool_message.tool_call_id}): "
                    f"token count ({tool_token_cnt}) is within threshold ({max_tool_message_tokens})",
                )
                continue

            # Save original full content before modifying
            original_content = tool_message.content

            # Generate file name from tool_call_id or create a unique identifier
            file_name = tool_message.tool_call_id or uuid4().hex
            store_path = store_dir / f"{file_name}.txt"

            # Store the full content for batch writing
            write_file_dict[store_path.as_posix()] = original_content

            # Log the compaction action
            logger.info(
                f"Compacting tool message (tool_call_id={tool_message.tool_call_id}): "
                f"token count={tool_token_cnt}, saving full content to {store_path.as_posix()}",
            )

            # Update tool message content with preview and file reference
            compact_result = f"tool call={file_name} result is stored in file path=`{store_path.as_posix()}`"
            if preview_char_length > 0:
                compact_result += f"\npreview: {original_content[:preview_char_length]}..."
            tool_message.content = compact_result

        # Store write_file_dict in context for potential batch writing
        self.context.write_file_dict = write_file_dict

        # Return the compacted messages as JSON
        self.context.response.answer = [x.simple_dump() for x in messages]
        self.context.response.metadata["write_file_dict"] = write_file_dict

        logger.info(f"Context compaction completed: {len(write_file_dict)} tool messages were compacted")

    async def async_default_execute(self, e: Exception = None, **_kwargs):
        """Handle execution errors by returning original messages.

        This method is called when an exception occurs during async_execute. It preserves
        the original messages and marks the operation as unsuccessful.

        Args:
            e: The exception that occurred during execution, if any.
            **_kwargs: Additional keyword arguments (unused but required by interface).
        """
        self.context.response.answer = self.context.messages
        self.context.response.success = False
        self.context.response.metadata["error"] = str(e)