Add token-based compaction as an alternative to invocation count based compaction

[jhwanflow]

Problem Statement

While using the current invocation-based compaction in Google ADK, I've encountered two significant pain points:

1. Context Length Overflow Before Reaching Invocation Threshold
   The context length can exceed the model's limit before the invocation count reaches the compaction threshold. Since compaction only triggers after a fixed number of invocations, there's no protection against context overflow when individual invocations contain large amounts of content.

2. Unnecessary Compaction Triggers
   Compaction is triggered even when the actual token count is relatively small, simply because the invocation count reached the threshold.

Summary

• Better alignment with model context limits (e.g., Gemini's 1M/2M windows)
• More efficient context usage for variable-length conversations
• Predictable behavior regardless of message length variations
• Cost optimization by compacting only when necessary

Suggestion

class TokenBasedCompactionConfig(EventsCompactionConfig):
    """configuration for token based compaction
    - token_threshold: Trigger compaction when input tokens exceed this value
    """
    token_threshold: int = 150000
    compaction_interval : Optional[int] = None


def _count_tokens_in_events(events: list[Event],
                                        tokenizer,
                                        ) -> int:
    total_tokens = 0
    for event in events:
        if event.content and event.content.parts:
            for part in event.content.parts:
                if part.text:
                    tokens = tokenzier.encode(part.text)
                    total_tokens += len(tokens)

    return total_tokens


async def _run_compaction_for_sliding_window_by_token(
    app: App,
    session: Session, 
    session_service: BaseSessionService,
):
    events = session.events
    if not events:
        return None
    
    last_compacted_end_timestamp = 0.0
    for event in reversed(events):
        if (
            event.actions
            and event.actions.compaction
            and event.actions.compaction.end_timestamp
        ):
            last_compacted_end_timestamp = event.actions.compaction.end_timestamp
            break

    new_events = [
        e for e in events
        if e.timestamp > last_compacted_end_timestamp
        and not (e.actions and e.actions.compaction)
    ]   

    if not new_events:
        return None

    # Count the number of tokens in new events
    total_tokens = _count_tokens_in_events(new_events)

    token_threshold = app.events_compaction_config.token_threshold
    
    logger.debug(f'New events token count: {total_tokens}, threshold: {token_threshold}')
    
    if total_tokens < token_threshold:
        return None

    
    invocation_latest_timestamps = {}
    for event in events:
        if event.invocation_id and not (event.actions and event.actions.compaction):
            invocation_latest_timestamps[event.invocation_id] = max(
                invocation_latest_timestamps.get(event.invocation_id, 0.0),
                event.timestamp,
            )

    unique_invocation_ids = list(invocation_latest_timestamps.keys())

    new_invocation_ids = [
        inv_id
        for inv_id in unique_invocation_ids
        if invocation_latest_timestamps[inv_id] > last_compacted_end_timestamp
    ]

    if not new_invocation_ids:
        return None

    end_inv_id = new_invocation_ids[-1]
    first_new_inv_id = new_invocation_ids[0]
    first_new_inv_idx = unique_invocation_ids.index(first_new_inv_id)

    overlap_size = getattr(
        app.events_compaction_config, 'overlap_size', 1
    )
    start_idx = max(0, first_new_inv_idx - overlap_size)
    start_inv_id = unique_invocation_ids[start_idx]

    last_event_idx = -1
    for i in range(len(events) - 1, -1, -1):
        if events[i].invocation_id == end_inv_id:
            last_event_idx = i
            break

    events_to_compact = []
    if last_event_idx != -1:
        first_event_start_inv_idx = -1
        for i, event in enumerate(events):
            if event.invocation_id == start_inv_id:
                first_event_start_inv_idx = i
                break
        if first_event_start_inv_idx != -1:
            events_to_compact = events[first_event_start_inv_idx : last_event_idx + 1]
            events_to_compact = [
                e
                for e in events_to_compact
                if not (e.actions and e.actions.compaction)
            ]

    if not events_to_compact:
        return None

    if not app.events_compaction_config.summarizer:
        app.events_compaction_config.summarizer = LlmEventSummarizer(
            llm= app.root_agent.canonical_model,
            prompt_template=summarize_prompt,
        )

    compaction_event = (
        await app.events_compaction_config.summarizer.maybe_summarize_events(
            events=events_to_compact,

        )
    )
    
    if compaction_event:
        await session_service.append_event(session=session, event=compaction_event)
        logger.info(
            f'Compaction completed. Summarized {len(events_to_compact)} events '
            f'({total_tokens} tokens) into 1 compacted event.'
        )
    logger.debug('Event compactor finished.')

[wdai-chwy]

+1