ReMe/reme2/memory/summarizer.py

"""Summarizer module for memory summarization operations."""

import datetime
import json
import zoneinfo

from agentscope.agent import ReActAgent
from agentscope.message import Msg
from agentscope.token import HuggingFaceTokenCounter
from agentscope.tool import Toolkit
from loguru import logger

from ..component import BaseStep
from ..schema import AsMsgStat, AsBlockStat


class Summarizer(BaseStep):
    """Summarizer step for summarizing memory messages."""

    def __init__(
        self,
        working_dir: str,
        memory_dir: str,
        memory_compact_threshold: int,
        toolkit: Toolkit | None = None,
        console_enabled: bool = False,
        timezone: str | None = None,
        add_thinking_block: bool = True,
        as_token_counter: HuggingFaceTokenCounter | None = None,
        **kwargs,
    ):
        """Initialize the summarizer step.

        Args:
            working_dir: Working directory path.
            memory_dir: Memory directory path for storing summaries.
            memory_compact_threshold: Token threshold for memory compaction.
            toolkit: Optional toolkit for the agent.
            console_enabled: Whether to enable console output.
            timezone: Optional timezone string for date formatting.
            add_thinking_block: Whether to include thinking blocks in output.
            as_token_counter: Optional token counter instance.
            **kwargs: Additional keyword arguments passed to BaseStep.
        """
        super().__init__(**kwargs)
        self.working_dir: str = working_dir
        self.memory_dir: str = memory_dir
        self.memory_compact_threshold: int = memory_compact_threshold
        self.toolkit: Toolkit | None = toolkit
        self.console_enabled: bool = console_enabled
        self.timezone: str | None = timezone
        self.add_thinking_block: bool = add_thinking_block
        self._as_token_counter: HuggingFaceTokenCounter | None = as_token_counter

    def _get_current_datetime(self) -> datetime.datetime:
        """Get current datetime with timezone, fallback to local time if timezone is invalid."""
        if self.timezone:
            try:
                return datetime.datetime.now(zoneinfo.ZoneInfo(self.timezone))
            except Exception as e:
                self.logger.error(f"Invalid timezone: {self.timezone}, falling back to local time error={e}")
        return datetime.datetime.now()

    async def _count_str_token(self, text: str) -> int:
        """Count tokens in a string."""
        return await self.as_token_counter.count(messages=[], text=text)

    async def _format_tool_result_output(self, output: str | list[dict]) -> tuple[str, int]:
        """Convert tool result output to string."""
        if isinstance(output, str):
            return output, await self._count_str_token(output)

        textual_parts = []
        total_token_count = 0
        for block in output:
            try:
                if not isinstance(block, dict) or "type" not in block:
                    logger.warning(
                        f"Invalid block: {block}, expected a dict with 'type' key, skipped.",
                    )
                    continue

                block_type = block["type"]

                if block_type == "text":
                    textual_parts.append(block.get("text", ""))
                    total_token_count += await self._count_str_token(textual_parts[-1])

                elif block_type in ["image", "audio", "video"]:
                    source = block.get("source", {})
                    if source.get("type") == "base64":
                        data = source.get("data", "")
                        total_token_count += len(data) // 4 if data else 10
                    else:
                        url = source.get("url", "")
                        total_token_count += await self._count_str_token(url) if url else 10
                        textual_parts.append(f"[{block_type}] {url}")

                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}")
                    total_token_count += await self._count_str_token(file_path)

                else:
                    logger.warning(
                        f"Unsupported block type '{block_type}' in tool result, skipped.",
                    )

            except Exception as e:
                logger.warning(
                    f"Failed to process block {block}: {e}, skipped.",
                )

        return "\n".join(textual_parts), total_token_count

    async def _stat_message(self, message: Msg) -> AsMsgStat:
        """Analyze a message and generate block statistics."""
        blocks = []
        if isinstance(message.content, str):
            blocks.append(
                AsBlockStat(
                    block_type="text",
                    text=message.content,
                    token_count=await self._count_str_token(message.content),
                ),
            )
            return AsMsgStat(
                name=message.name or message.role,
                role=message.role,
                content=blocks,
                timestamp=message.timestamp or "",
                metadata=message.metadata or {},
            )

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

            if block_type == "text":
                text = block.get("text", "")
                token_count = await 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 = await 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", "")
                if source.get("type") == "base64":
                    data = source.get("data", "")
                    token_count = len(data) // 4 if data else 10
                else:
                    token_count = await 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 = await 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, token_count = await self._format_tool_result_output(output)
                blocks.append(
                    AsBlockStat(
                        block_type=block_type,
                        text="",
                        token_count=token_count,
                        tool_name=tool_name,
                        tool_output=formatted_output,
                    ),
                )

            else:
                logger.warning(f"Unsupported block type {block_type}, skipped.")

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

    async def _count_msgs_token(self, messages: list[Msg]) -> int:
        """Count total token count of a list of messages."""
        total = 0
        for msg in messages:
            stat = await self._stat_message(msg)
            total += stat.total_tokens
        return total

    async def _format_msgs_to_str(
        self,
        messages: list[Msg],
        memory_compact_threshold: int,
        include_thinking: bool = True,
    ) -> 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 = await self._stat_message(messages[i])
            formatted_content = stat.format(include_thinking=include_thinking)
            content_token_count = await self._count_str_token(formatted_content)

            is_latest = i == len(messages) - 1
            if not is_latest and total_token_count + content_token_count > memory_compact_threshold:
                logger.info(
                    f"Skipping older messages: adding {content_token_count} tokens would exceed threshold "
                    f"{memory_compact_threshold} (current: {total_token_count})",
                )
                break

            if is_latest and content_token_count > memory_compact_threshold:
                logger.warning(
                    f"Latest message alone ({content_token_count} tokens) exceeds threshold "
                    f"{memory_compact_threshold}, including it anyway.",
                )

            formatted_parts.append(formatted_content)
            total_token_count += content_token_count

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

    async def execute(self):
        """Execute the summarization step."""
        messages: list[Msg] = self.context.data.get("messages", [])

        if not messages:
            return ""

        before_token_count = await self._count_msgs_token(messages)
        history_formatted_str: str = await self._format_msgs_to_str(
            messages=messages,
            memory_compact_threshold=self.memory_compact_threshold,
            include_thinking=self.add_thinking_block,
        )
        after_token_count = await self._count_str_token(history_formatted_str)
        logger.info(f"Summarizer before_token_count={before_token_count} after_token_count={after_token_count}")

        if not history_formatted_str:
            logger.warning(f"No history to summarize. messages={messages}")
            return ""

        agent = ReActAgent(
            name="reme_summarizer",
            model=self.as_llm.model,
            sys_prompt="You are a helpful assistant.",
            formatter=self.as_llm_formatter.formatter,
            toolkit=self.toolkit,
        )
        agent.set_console_output_enabled(self.console_enabled)

        user_message: str = f"# conversation\n{history_formatted_str}\n\n" + self.prompt_format(
            "user_message",
            date=self._get_current_datetime().strftime("%Y-%m-%d"),
            working_dir=self.working_dir,
            memory_dir=self.memory_dir,
        )

        summary_msg: Msg = await agent.reply(
            Msg(
                name="reme",
                role="user",
                content=user_message,
            ),
        )
        for i, (msg, _) in enumerate(agent.memory.content):
            logger.info(f"Summarizer memory[{i}]: {msg.content}")

        history_summary: str = summary_msg.get_text_content()
        logger.info(f"Summarizer Result:\n{history_summary}")
        return history_summary