ReMe/reme/memory/file_based/as_msg_handler.py

"""Handler for AgentScope message processing, token counting, and context management."""

import json

from agentscope.message import Msg
from agentscope.token import HuggingFaceTokenCounter

from ...core.schema import AsMsgStat, AsBlockStat
from ...core.utils import get_std_logger

logger = get_std_logger()


class AsMsgHandler:
    """Handles token counting, formatting, and context compaction for AgentScope messages."""

    def __init__(self, token_counter: HuggingFaceTokenCounter):
        self._token_counter = token_counter

    def count_str_token(self, text: str) -> int:
        """Count tokens in a string.

        Args:
            text: The text to count tokens for.

        Returns:
            The number of tokens in the text.
        """
        if not text:
            return 0

        try:
            token_ids = self._token_counter.tokenizer.encode(text)
            token_count = len(token_ids)
            return token_count

        except Exception as e:
            estimated_tokens = len(text.encode("utf-8")) // 4
            logger.warning(f"Failed to count string tokens: {text}, e={e}")
            return estimated_tokens

    @staticmethod
    def _format_tool_result_output(output: str | list[dict]) -> str:
        """Convert tool result output to string.

        Args:
            output: Tool result output, either string or list of content blocks.

        Returns:
            Formatted string representation of the tool result.
        """
        if isinstance(output, str):
            return output

        textual_parts = []

        for block in output:
            try:
                if not isinstance(block, dict) or "type" not in block:
                    logger.warning(
                        "Invalid block: %s, expected a dict with 'type' key, skipped.",
                        block,
                    )
                    continue

                block_type = block["type"]

                if block_type == "text":
                    textual_parts.append(block.get("text", ""))

                elif block_type in ["image", "audio", "video"]:
                    source = block.get("source", {})
                    url = source.get("url", "")
                    if url:
                        textual_parts.append(f"[{block_type}] {url}")
                    else:
                        textual_parts.append(f"[{block_type}]")

                elif block_type == "file":
                    file_path = block.get("path", "") or block.get("url", "")
                    file_name = block.get("name", file_path)
                    textual_parts.append(f"[file] {file_name}: {file_path}")

                else:
                    logger.warning(
                        "Unsupported block type '%s' in tool result, skipped.",
                        block_type,
                    )

            except Exception as e:
                logger.warning(
                    "Failed to process block %s: %s, skipped.",
                    block,
                    e,
                )

        if not textual_parts:
            return ""
        if len(textual_parts) == 1:
            return textual_parts[0]
        return "\n".join(f"- {part}" for part in textual_parts)

    def stat_message(self, message: Msg) -> AsMsgStat:
        """Analyze a message and generate block statistics."""
        blocks = []

        for block in message.get_content_blocks():
            block_type = block.get("type", "unknown")

            if block_type == "text":
                text = block.get("text", "")
                token_count = self.count_str_token(text)
                blocks.append(
                    AsBlockStat(
                        block_type=block_type,
                        text=text,
                        token_count=token_count,
                    ),
                )

            elif block_type == "thinking":
                thinking = block.get("thinking", "")
                token_count = self.count_str_token(thinking)
                blocks.append(
                    AsBlockStat(
                        block_type=block_type,
                        text=thinking,
                        token_count=token_count,
                    ),
                )

            elif block_type in ("image", "audio", "video"):
                source = block.get("source", {})
                url = source.get("url", "")
                # For media, estimate fixed token cost or count URL
                if source.get("type") == "base64":
                    data = source.get("data", "")
                    token_count = len(data) // 4 if data else 10
                else:
                    token_count = self.count_str_token(url) if url else 10
                blocks.append(
                    AsBlockStat(
                        block_type=block_type,
                        text="",
                        token_count=token_count,
                        media_url=url,
                    ),
                )

            elif block_type == "tool_use":
                tool_name = block.get("name", "")
                tool_input = block.get("input", {})
                try:
                    input_str = json.dumps(tool_input, ensure_ascii=False)
                except (TypeError, ValueError):
                    input_str = str(tool_input)
                token_count = self.count_str_token(tool_name + input_str)
                blocks.append(
                    AsBlockStat(
                        block_type=block_type,
                        text="",
                        token_count=token_count,
                        tool_name=tool_name,
                        tool_input=input_str,
                    ),
                )

            elif block_type == "tool_result":
                tool_name = block.get("name", "")
                output = block.get("output", "")
                formatted_output = self._format_tool_result_output(output)
                token_count = self.count_str_token(formatted_output)
                blocks.append(
                    AsBlockStat(
                        block_type=block_type,
                        text="",
                        token_count=token_count,
                        tool_name=tool_name,
                        tool_output=formatted_output,
                    ),
                )

            else:
                logger.warning("Unsupported block type %s, skipped.", block_type)

        return AsMsgStat(
            name=message.name or message.role,
            role=message.role,
            content=blocks,
            timestamp=message.timestamp or "",
            metadata=message.metadata or {},
        )

    def format_msgs_to_str(
        self,
        messages: list[Msg],
        memory_compact_threshold: int,
        include_thinking: bool = False,
    ) -> str:
        """Format list of messages to a single formatted string.

        Messages are processed in reverse order (newest first) and older
        messages are skipped when token count exceeds memory_compact_threshold.

        Args:
            messages: List of Msg objects to format.
            memory_compact_threshold: Maximum token count before skipping older messages.
            include_thinking: Whether to include thinking blocks in output.
        """
        if not messages:
            return ""

        formatted_parts: list[str] = []
        total_token_count = 0

        for i in range(len(messages) - 1, -1, -1):
            stat = self.stat_message(messages[i])

            if total_token_count + stat.total_tokens > memory_compact_threshold:
                logger.info(
                    "Skipping older messages: adding %d tokens would exceed threshold %d (current: %d)",
                    stat.total_tokens,
                    memory_compact_threshold,
                    total_token_count,
                )
                break

            formatted_parts.append(stat.format(include_thinking=include_thinking))
            total_token_count += stat.total_tokens

        formatted_parts.reverse()
        return "\n\n".join(formatted_parts)

    def context_check(
        self,
        messages: list[Msg],
        memory_compact_threshold: int,
        memory_compact_reserve: int,
    ) -> tuple[list[Msg], list[Msg]]:
        """Check if context exceeds threshold and split messages accordingly.

        This method checks if the total token count of messages exceeds the
        memory_compact_threshold. If not, returns empty list and original messages.
        If exceeded, uses memory_compact_reserve as the limit to keep messages
        from the end, ensuring tool_use and tool_result blocks are properly paired.

        Args:
            messages: List of Msg objects to check.
            memory_compact_threshold: Maximum token count threshold to trigger compaction.
            memory_compact_reserve: Token limit for messages to keep after compaction.

        Returns:
            A tuple of (messages_to_compact, messages_to_keep):
            - messages_to_compact: Older messages that need to be compacted
            - messages_to_keep: Recent messages within the reserve limit
        """
        if not messages:
            return [], []

        # Calculate total tokens and stats for all messages
        msg_stats: list[tuple[Msg, AsMsgStat]] = []
        total_tokens = 0
        for msg in messages:
            stat = self.stat_message(msg)
            msg_stats.append((msg, stat))
            total_tokens += stat.total_tokens

        # If total tokens don't exceed threshold, no compaction needed
        if total_tokens <= memory_compact_threshold:
            return [], messages

        # Collect all tool_use ids and their message indices
        # tool_use_id -> message index
        tool_use_locations: dict[str, int] = {}
        # tool_result_id -> message index
        tool_result_locations: dict[str, int] = {}

        for idx, (msg, _) in enumerate(msg_stats):
            for block in msg.get_content_blocks("tool_use"):
                tool_id = block.get("id", "")
                if tool_id:
                    tool_use_locations[tool_id] = idx

            for block in msg.get_content_blocks("tool_result"):
                tool_id = block.get("id", "")
                if tool_id:
                    tool_result_locations[tool_id] = idx

        # Iterate from the end, accumulating messages to keep within reserve limit
        keep_indices: set[int] = set()
        accumulated_tokens = 0

        for i in range(len(msg_stats) - 1, -1, -1):
            msg, stat = msg_stats[i]

            # Check if adding this message would exceed reserve limit
            if accumulated_tokens + stat.total_tokens > memory_compact_reserve:
                logger.info(
                    "Context check: adding message %d with %d tokens would exceed reserve %d (current: %d)",
                    i,
                    stat.total_tokens,
                    memory_compact_reserve,
                    accumulated_tokens,
                )
                break

            # Check tool_result dependencies - if this message has tool_result,
            # we need to ensure the corresponding tool_use is also included
            tool_result_ids = [
                block.get("id", "") for block in msg.get_content_blocks("tool_result") if block.get("id", "")
            ]

            # Calculate extra tokens needed for dependent tool_use messages
            extra_tokens = 0
            dependent_indices: set[int] = set()

            for tool_id in tool_result_ids:
                if tool_id in tool_use_locations:
                    tool_use_idx = tool_use_locations[tool_id]
                    if tool_use_idx not in keep_indices and tool_use_idx != i:
                        dependent_indices.add(tool_use_idx)
                        _, dep_stat = msg_stats[tool_use_idx]
                        extra_tokens += dep_stat.total_tokens

            # Check if we can fit this message plus its dependencies within reserve
            if accumulated_tokens + stat.total_tokens + extra_tokens > memory_compact_reserve:
                logger.info(
                    "Context check: message %d requires %d extra tokens for tool_use dependencies, "
                    "total would exceed reserve %d",
                    i,
                    extra_tokens,
                    memory_compact_reserve,
                )
                break

            # Add this message and its dependencies
            keep_indices.add(i)
            keep_indices.update(dependent_indices)
            accumulated_tokens += stat.total_tokens + extra_tokens

        # Build final lists based on keep_indices (preserve original order)
        messages_to_compact = []
        messages_to_keep = []

        for idx, (msg, _) in enumerate(msg_stats):
            if idx in keep_indices:
                messages_to_keep.append(msg)
            else:
                messages_to_compact.append(msg)

        logger.info(
            "Context check result: %d messages to compact, %d messages to keep, "
            "total tokens: %d, threshold: %d, reserve: %d, kept tokens: %d",
            len(messages_to_compact),
            len(messages_to_keep),
            total_tokens,
            memory_compact_threshold,
            memory_compact_reserve,
            accumulated_tokens,
        )

        return messages_to_compact, messages_to_keep