_summarizer_checker.py

# Tencent is pleased to support the open source community by making tRPC-Agent-Python available.
#
# Copyright (C) 2026 Tencent. All rights reserved.
#
# tRPC-Agent-Python is licensed under Apache-2.0.
"""Summarizer checker functions."""

from __future__ import annotations

import time
from typing import Callable
from typing import List

from trpc_agent_sdk.events import Event
from trpc_agent_sdk.log import logger

from ._session import Session

CheckSummarizerFunction = Callable[[Session], bool]


def _leading_summary_event(session: Session) -> Event | None:
    """Return the summary anchor when the active event window starts with one."""
    if session.events and session.events[0].is_summary_event():
        return session.events[0]
    return None


def _events_after_summary_anchor(session: Session) -> list[Event]:
    """Return active events after the leading summary anchor."""
    return session.events[1:] if _leading_summary_event(session) else session.events


def set_summarizer_token_threshold(token_count: int) -> CheckSummarizerFunction:
    """Set the token threshold for summarizer.

    Args:
        token_count: The token count to check

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        # Filter events with usage_metadata
        events_with_metadata = [
            event for event in _events_after_summary_anchor(session) if event.usage_metadata is not None
        ]

        # If no events have usage_metadata, log a warning and return False
        if not events_with_metadata:
            logger.warning(
                "No events with usage_metadata found in session %s. "
                "Token-based summarization check returning False. "
                "This may indicate that LLM responses are not properly recording token usage.", session.id)
            return False

        # Calculate total token count
        total_token_count = sum(event.usage_metadata.total_token_count for event in events_with_metadata)

        return total_token_count > token_count

    return _decorator


def set_summarizer_events_count_threshold(event_count: int = 30) -> CheckSummarizerFunction:
    """Set the number of events threshold for summarizer.

    Args:
        event_count: The event count to check

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        # Check if we have enough events to warrant summarization
        return len(_events_after_summary_anchor(session)) > event_count

    return _decorator


def set_summarizer_time_interval_threshold(time_interval: float = 300.0) -> CheckSummarizerFunction:
    """Set the time interval threshold for summarizer.

    Args:
        time_interval: The time interval to check

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        summary_event = _leading_summary_event(session)
        events_after_summary = _events_after_summary_anchor(session)
        if not events_after_summary:
            return False

        if summary_event is not None:
            return time.time() - summary_event.timestamp > time_interval

        return time.time() - events_after_summary[-1].timestamp > time_interval

    return _decorator


def set_summarizer_important_content_threshold(important_content_count: int = 10) -> CheckSummarizerFunction:
    """Set the important content threshold for summarizer.

    Args:
        important_content: The important content to check

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        # Check if there's important content to summarize
        for event in _events_after_summary_anchor(session):
            if event.content and event.content.parts:
                for part in event.content.parts:
                    if part.text and len(part.text.strip()) > important_content_count:
                        return True
        return False

    return _decorator


def set_summarizer_conversation_threshold(conversation_count: int = 100) -> CheckSummarizerFunction:
    """Set the conversation count threshold for summarizer.

    Args:
        conversation_count: The conversation count to check

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        if session.conversation_count > conversation_count:
            session.conversation_count = 0
            return True
        return False

    return _decorator


def set_summarizer_check_functions_by_and(funcs: List[CheckSummarizerFunction]) -> CheckSummarizerFunction:
    """Set the check functions for summarizer, all the functions must return True.
    Args:
        funcs: The list of check summarizer functions

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        return all(func(session) for func in funcs)

    return _decorator


def set_summarizer_check_functions_by_or(funcs: List[CheckSummarizerFunction]) -> CheckSummarizerFunction:
    """Set the check functions for summarizer, any of the functions return True.

    Args:
        funcs: The list of check summarizer functions

    Returns:
        True if summarization is needed, False otherwise
    """

    def _decorator(session: Session) -> bool:
        return any(func(session) for func in funcs)

    return _decorator