mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Add post-invocation token-threshold compaction with event retention
Adds optional token_limit and event_retention_size fields to EventsCompactionConfig. When the latest prompt token count meets/exceeds the threshold, ADK compacts older raw events after the invocation and keeps the last N events un-compacted. Updates prompt history building to apply compaction ranges correctly so retained events remain visible. Close #4146 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 868297968
This commit is contained in:
George Weale
committed by
Copybara-Service
Copybara-Service
parent
25ec2c6b61
commit
a88e864755
@@ -80,6 +80,33 @@ class EventsCompactionConfig(BaseModel):
|
||||
end of the last compacted range. This creates an overlap between consecutive
|
||||
compacted summaries, maintaining context."""
|
||||
|
||||
token_threshold: Optional[int] = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
)
|
||||
"""Post-invocation token threshold trigger.
|
||||
|
||||
If set, ADK will attempt a post-invocation compaction when the most recently
|
||||
observed prompt token count meets or exceeds this threshold.
|
||||
"""
|
||||
|
||||
event_retention_size: Optional[int] = Field(default=None, ge=0)
|
||||
"""Post-invocation raw event retention size.
|
||||
|
||||
If token-based post-invocation compaction is triggered, this keeps the last N
|
||||
raw events un-compacted.
|
||||
"""
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_token_params(self) -> EventsCompactionConfig:
|
||||
token_threshold_set = self.token_threshold is not None
|
||||
retention_size_set = self.event_retention_size is not None
|
||||
if token_threshold_set != retention_size_set:
|
||||
raise ValueError(
|
||||
"token_threshold and event_retention_size must be set together."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class App(BaseModel):
|
||||
"""Represents an LLM-backed agentic application.
|
||||
|
||||
@@ -16,14 +16,236 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from google.adk.apps.app import App
|
||||
from google.adk.apps.llm_event_summarizer import LlmEventSummarizer
|
||||
from google.adk.sessions.base_session_service import BaseSessionService
|
||||
from google.adk.sessions.session import Session
|
||||
from ..events.event import Event
|
||||
from ..sessions.base_session_service import BaseSessionService
|
||||
from ..sessions.session import Session
|
||||
from .app import App
|
||||
from .llm_event_summarizer import LlmEventSummarizer
|
||||
|
||||
logger = logging.getLogger('google_adk.' + __name__)
|
||||
|
||||
|
||||
def _count_text_chars_in_event(event: Event) -> int:
|
||||
"""Returns the number of text characters in an event's content."""
|
||||
total_chars = 0
|
||||
if event.content and event.content.parts:
|
||||
for part in event.content.parts:
|
||||
if part.text:
|
||||
total_chars += len(part.text)
|
||||
return total_chars
|
||||
|
||||
|
||||
def _is_compaction_subsumed(
|
||||
*,
|
||||
start_timestamp: float,
|
||||
end_timestamp: float,
|
||||
event_index: int,
|
||||
compactions: list[tuple[int, float, float, Event]],
|
||||
) -> bool:
|
||||
"""Returns True if a compaction range is fully contained by another.
|
||||
|
||||
If two compactions have identical ranges, the earlier event is treated as
|
||||
subsumed by the later event.
|
||||
"""
|
||||
for other_index, other_start, other_end, _ in compactions:
|
||||
if other_index == event_index:
|
||||
continue
|
||||
if other_start <= start_timestamp and other_end >= end_timestamp:
|
||||
if (
|
||||
other_start < start_timestamp
|
||||
or other_end > end_timestamp
|
||||
or other_index > event_index
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _estimate_prompt_token_count(events: list[Event]) -> int | None:
|
||||
"""Returns an approximate prompt token count from session events.
|
||||
|
||||
This estimate is compaction-aware: it counts compaction summaries and only
|
||||
counts raw events that would remain visible after applying compaction ranges.
|
||||
"""
|
||||
compactions: list[tuple[int, float, float, Event]] = []
|
||||
for i, event in enumerate(events):
|
||||
if not (event.actions and event.actions.compaction):
|
||||
continue
|
||||
compaction = event.actions.compaction
|
||||
if (
|
||||
compaction.start_timestamp is None
|
||||
or compaction.end_timestamp is None
|
||||
or compaction.compacted_content is None
|
||||
):
|
||||
continue
|
||||
compactions.append((
|
||||
i,
|
||||
compaction.start_timestamp,
|
||||
compaction.end_timestamp,
|
||||
Event(
|
||||
timestamp=compaction.end_timestamp,
|
||||
author='model',
|
||||
content=compaction.compacted_content,
|
||||
branch=event.branch,
|
||||
invocation_id=event.invocation_id,
|
||||
actions=event.actions,
|
||||
),
|
||||
))
|
||||
|
||||
effective_compactions = [
|
||||
(i, start, end, summary_event)
|
||||
for i, start, end, summary_event in compactions
|
||||
if not _is_compaction_subsumed(
|
||||
start_timestamp=start,
|
||||
end_timestamp=end,
|
||||
event_index=i,
|
||||
compactions=compactions,
|
||||
)
|
||||
]
|
||||
compaction_ranges = [
|
||||
(start, end) for _, start, end, _ in effective_compactions
|
||||
]
|
||||
|
||||
def _is_timestamp_compacted(ts: float) -> bool:
|
||||
for start_ts, end_ts in compaction_ranges:
|
||||
if start_ts <= ts <= end_ts:
|
||||
return True
|
||||
return False
|
||||
|
||||
total_chars = 0
|
||||
for _, _, _, summary_event in effective_compactions:
|
||||
total_chars += _count_text_chars_in_event(summary_event)
|
||||
|
||||
for event in events:
|
||||
if event.actions and event.actions.compaction:
|
||||
continue
|
||||
if _is_timestamp_compacted(event.timestamp):
|
||||
continue
|
||||
total_chars += _count_text_chars_in_event(event)
|
||||
|
||||
if total_chars <= 0:
|
||||
return None
|
||||
|
||||
# Rough estimate: 4 characters per token.
|
||||
return total_chars // 4
|
||||
|
||||
|
||||
def _latest_prompt_token_count(events: list[Event]) -> int | None:
|
||||
"""Returns the most recently observed prompt token count, if available."""
|
||||
for event in reversed(events):
|
||||
if (
|
||||
event.usage_metadata
|
||||
and event.usage_metadata.prompt_token_count is not None
|
||||
):
|
||||
return event.usage_metadata.prompt_token_count
|
||||
return _estimate_prompt_token_count(events)
|
||||
|
||||
|
||||
def _latest_compaction_event(events: list[Event]) -> Event | None:
|
||||
"""Returns the compaction event with the greatest covered end timestamp."""
|
||||
latest_event = None
|
||||
latest_end = 0.0
|
||||
for event in events:
|
||||
if (
|
||||
event.actions
|
||||
and event.actions.compaction
|
||||
and event.actions.compaction.end_timestamp is not None
|
||||
):
|
||||
end_ts = event.actions.compaction.end_timestamp
|
||||
if end_ts is not None and end_ts >= latest_end:
|
||||
latest_end = end_ts
|
||||
latest_event = event
|
||||
return latest_event
|
||||
|
||||
|
||||
def _latest_compaction_end_timestamp(events: list[Event]) -> float:
|
||||
"""Returns the end timestamp of the most recent compaction event."""
|
||||
latest_event = _latest_compaction_event(events)
|
||||
if not latest_event or not latest_event.actions.compaction:
|
||||
return 0.0
|
||||
if latest_event.actions.compaction.end_timestamp is None:
|
||||
return 0.0
|
||||
return latest_event.actions.compaction.end_timestamp
|
||||
|
||||
|
||||
async def _run_compaction_for_token_threshold(
|
||||
app: App, session: Session, session_service: BaseSessionService
|
||||
):
|
||||
"""Runs post-invocation compaction based on a token threshold.
|
||||
|
||||
If triggered, this compacts older raw events and keeps the last
|
||||
`event_retention_size` raw events un-compacted.
|
||||
"""
|
||||
config = app.events_compaction_config
|
||||
if not config:
|
||||
return False
|
||||
if config.token_threshold is None or config.event_retention_size is None:
|
||||
return False
|
||||
|
||||
prompt_token_count = _latest_prompt_token_count(session.events)
|
||||
if prompt_token_count is None or prompt_token_count < config.token_threshold:
|
||||
return False
|
||||
|
||||
latest_compaction_event = _latest_compaction_event(session.events)
|
||||
last_compacted_end_timestamp = 0.0
|
||||
if (
|
||||
latest_compaction_event
|
||||
and latest_compaction_event.actions
|
||||
and latest_compaction_event.actions.compaction
|
||||
and latest_compaction_event.actions.compaction.end_timestamp is not None
|
||||
):
|
||||
last_compacted_end_timestamp = (
|
||||
latest_compaction_event.actions.compaction.end_timestamp
|
||||
)
|
||||
candidate_events = [
|
||||
e
|
||||
for e in session.events
|
||||
if not (e.actions and e.actions.compaction)
|
||||
and e.timestamp > last_compacted_end_timestamp
|
||||
]
|
||||
|
||||
if len(candidate_events) <= config.event_retention_size:
|
||||
return False
|
||||
|
||||
if config.event_retention_size == 0:
|
||||
events_to_compact = candidate_events
|
||||
else:
|
||||
events_to_compact = candidate_events[: -config.event_retention_size]
|
||||
if not events_to_compact:
|
||||
return False
|
||||
|
||||
# Rolling summary: if a previous compaction exists, seed the next summary with
|
||||
# the previous compaction summary content so new compactions can subsume older
|
||||
# ones while still keeping `event_retention_size` raw events visible.
|
||||
if (
|
||||
latest_compaction_event
|
||||
and latest_compaction_event.actions
|
||||
and latest_compaction_event.actions.compaction
|
||||
and latest_compaction_event.actions.compaction.start_timestamp is not None
|
||||
and latest_compaction_event.actions.compaction.compacted_content
|
||||
is not None
|
||||
):
|
||||
seed_event = Event(
|
||||
timestamp=latest_compaction_event.actions.compaction.start_timestamp,
|
||||
author='model',
|
||||
content=latest_compaction_event.actions.compaction.compacted_content,
|
||||
branch=latest_compaction_event.branch,
|
||||
invocation_id=Event.new_id(),
|
||||
)
|
||||
events_to_compact = [seed_event] + events_to_compact
|
||||
|
||||
if not config.summarizer:
|
||||
config.summarizer = LlmEventSummarizer(llm=app.root_agent.canonical_model)
|
||||
|
||||
compaction_event = await config.summarizer.maybe_summarize_events(
|
||||
events=events_to_compact
|
||||
)
|
||||
if compaction_event:
|
||||
await session_service.append_event(session=session, event=compaction_event)
|
||||
logger.debug('Token-threshold event compactor finished.')
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def _run_compaction_for_sliding_window(
|
||||
app: App, session: Session, session_service: BaseSessionService
|
||||
):
|
||||
@@ -109,6 +331,18 @@ async def _run_compaction_for_sliding_window(
|
||||
events = session.events
|
||||
if not events:
|
||||
return None
|
||||
|
||||
# Prefer token-threshold compaction if configured and triggered.
|
||||
if (
|
||||
app.events_compaction_config
|
||||
and app.events_compaction_config.token_threshold is not None
|
||||
):
|
||||
token_compacted = await _run_compaction_for_token_threshold(
|
||||
app, session, session_service
|
||||
)
|
||||
if token_compacted:
|
||||
return None
|
||||
|
||||
# Find the last compaction event and its range.
|
||||
last_compacted_end_timestamp = 0.0
|
||||
for event in reversed(events):
|
||||
|
||||
@@ -323,50 +323,88 @@ def _process_compaction_events(events: list[Event]) -> list[Event]:
|
||||
Returns:
|
||||
A list of events with compaction applied.
|
||||
"""
|
||||
# example of compaction events:
|
||||
# [event_1(timestamp=1), event_2(timestamp=2),
|
||||
# compaction_1(event_1, event_2, timestamp=3), event_3(timestamp=4),
|
||||
# compaction_2(event_2, event_3, timestamp=5), event_4(timestamp=6)]
|
||||
# for each compaction event, it only covers the events at most between the
|
||||
# current compaction and the previous compaction. So during compaction, we
|
||||
# don't have to go across compaction boundaries.
|
||||
# Compaction events are always strictly in order based on event timestamp.
|
||||
events_to_process = []
|
||||
last_compaction_start_time = float('inf')
|
||||
# Example:
|
||||
# [event_1(ts=1), event_2(ts=2), compaction_1(1-2), event_3(ts=4),
|
||||
# compaction_2(2-4), event_4(ts=6)].
|
||||
#
|
||||
# Overlaps are resolved by keeping only non-subsumed compaction summaries.
|
||||
# A summary event is materialized at its compaction end timestamp, and raw
|
||||
# events inside any kept compaction range are filtered out.
|
||||
compaction_infos: list[tuple[int, float, float]] = []
|
||||
for i, event in enumerate(events):
|
||||
if not (event.actions and event.actions.compaction):
|
||||
continue
|
||||
compaction = event.actions.compaction
|
||||
if (
|
||||
compaction.start_timestamp is None
|
||||
or compaction.end_timestamp is None
|
||||
or compaction.compacted_content is None
|
||||
):
|
||||
continue
|
||||
compaction_infos.append(
|
||||
(i, compaction.start_timestamp, compaction.end_timestamp)
|
||||
)
|
||||
|
||||
# Iterate in reverse to easily handle overlapping compactions.
|
||||
for event in reversed(events):
|
||||
subsumed_compaction_event_indexes: set[int] = set()
|
||||
for event_index, start_ts, end_ts in compaction_infos:
|
||||
for other_index, other_start, other_end in compaction_infos:
|
||||
if other_index == event_index:
|
||||
continue
|
||||
if other_start <= start_ts and other_end >= end_ts:
|
||||
if (
|
||||
other_start < start_ts
|
||||
or other_end > end_ts
|
||||
or other_index > event_index
|
||||
):
|
||||
subsumed_compaction_event_indexes.add(event_index)
|
||||
break
|
||||
|
||||
compaction_ranges: list[tuple[float, float]] = []
|
||||
processed_items: list[tuple[float, int, Event]] = []
|
||||
|
||||
for i, event in enumerate(events):
|
||||
if event.actions and event.actions.compaction:
|
||||
if i in subsumed_compaction_event_indexes:
|
||||
continue
|
||||
compaction = event.actions.compaction
|
||||
if (
|
||||
compaction.start_timestamp is not None
|
||||
and compaction.end_timestamp is not None
|
||||
compaction.start_timestamp is None
|
||||
or compaction.end_timestamp is None
|
||||
or compaction.compacted_content is None
|
||||
):
|
||||
# Create a new event for the compacted summary.
|
||||
new_event = Event(
|
||||
timestamp=compaction.end_timestamp,
|
||||
author='model',
|
||||
content=compaction.compacted_content,
|
||||
branch=event.branch,
|
||||
invocation_id=event.invocation_id,
|
||||
actions=event.actions,
|
||||
)
|
||||
# Prepend to maintain chronological order in the final list.
|
||||
events_to_process.insert(0, new_event)
|
||||
# Update the boundary for filtering. Events with timestamps greater than
|
||||
# or equal to this start time have been compacted.
|
||||
last_compaction_start_time = min(
|
||||
last_compaction_start_time, compaction.start_timestamp
|
||||
)
|
||||
elif event.timestamp < last_compaction_start_time:
|
||||
# This event is not a compaction and is before the current compaction
|
||||
# range. Prepend to maintain chronological order.
|
||||
events_to_process.insert(0, event)
|
||||
else:
|
||||
# skip the event
|
||||
pass
|
||||
continue
|
||||
compaction_ranges.append(
|
||||
(compaction.start_timestamp, compaction.end_timestamp)
|
||||
)
|
||||
processed_items.append((
|
||||
compaction.end_timestamp,
|
||||
i,
|
||||
Event(
|
||||
timestamp=compaction.end_timestamp,
|
||||
author='model',
|
||||
content=compaction.compacted_content,
|
||||
branch=event.branch,
|
||||
invocation_id=event.invocation_id,
|
||||
actions=event.actions,
|
||||
),
|
||||
))
|
||||
|
||||
return events_to_process
|
||||
def _is_timestamp_compacted(ts: float) -> bool:
|
||||
for start_ts, end_ts in compaction_ranges:
|
||||
if start_ts <= ts <= end_ts:
|
||||
return True
|
||||
return False
|
||||
|
||||
for i, event in enumerate(events):
|
||||
if event.actions and event.actions.compaction:
|
||||
continue
|
||||
if _is_timestamp_compacted(event.timestamp):
|
||||
continue
|
||||
processed_items.append((event.timestamp, i, event))
|
||||
|
||||
# Keep chronological order and a stable tie-breaker for equal timestamps.
|
||||
processed_items.sort(key=lambda item: (item[0], item[1]))
|
||||
return [event for _, _, event in processed_items]
|
||||
|
||||
|
||||
def _get_contents(
|
||||
|
||||
@@ -20,6 +20,7 @@ from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.apps.app import App
|
||||
from google.adk.apps.app import EventsCompactionConfig
|
||||
from google.adk.apps.compaction import _run_compaction_for_sliding_window
|
||||
import google.adk.apps.compaction as compaction_module
|
||||
from google.adk.apps.llm_event_summarizer import LlmEventSummarizer
|
||||
from google.adk.events.event import Event
|
||||
from google.adk.events.event_actions import EventActions
|
||||
@@ -27,8 +28,10 @@ from google.adk.events.event_actions import EventCompaction
|
||||
from google.adk.flows.llm_flows import contents
|
||||
from google.adk.sessions.base_session_service import BaseSessionService
|
||||
from google.adk.sessions.session import Session
|
||||
from google.genai import types
|
||||
from google.genai.types import Content
|
||||
from google.genai.types import Part
|
||||
from pydantic import ValidationError
|
||||
import pytest
|
||||
|
||||
|
||||
@@ -42,17 +45,31 @@ class TestCompaction(unittest.IsolatedAsyncioTestCase):
|
||||
self.mock_compactor = AsyncMock(spec=LlmEventSummarizer)
|
||||
|
||||
def _create_event(
|
||||
self, timestamp: float, invocation_id: str, text: str
|
||||
self,
|
||||
timestamp: float,
|
||||
invocation_id: str,
|
||||
text: str,
|
||||
prompt_token_count: int | None = None,
|
||||
) -> Event:
|
||||
usage_metadata = None
|
||||
if prompt_token_count is not None:
|
||||
usage_metadata = types.GenerateContentResponseUsageMetadata(
|
||||
prompt_token_count=prompt_token_count
|
||||
)
|
||||
return Event(
|
||||
timestamp=timestamp,
|
||||
invocation_id=invocation_id,
|
||||
author='user',
|
||||
content=Content(role='user', parts=[Part(text=text)]),
|
||||
usage_metadata=usage_metadata,
|
||||
)
|
||||
|
||||
def _create_compacted_event(
|
||||
self, start_ts: float, end_ts: float, summary_text: str
|
||||
self,
|
||||
start_ts: float,
|
||||
end_ts: float,
|
||||
summary_text: str,
|
||||
appended_ts: float | None = None,
|
||||
) -> Event:
|
||||
compaction = EventCompaction(
|
||||
start_timestamp=start_ts,
|
||||
@@ -62,7 +79,7 @@ class TestCompaction(unittest.IsolatedAsyncioTestCase):
|
||||
),
|
||||
)
|
||||
return Event(
|
||||
timestamp=end_ts,
|
||||
timestamp=appended_ts if appended_ts is not None else end_ts,
|
||||
author='compactor',
|
||||
content=compaction.compacted_content,
|
||||
actions=EventActions(compaction=compaction),
|
||||
@@ -225,6 +242,246 @@ class TestCompaction(unittest.IsolatedAsyncioTestCase):
|
||||
self.mock_compactor.maybe_summarize_events.assert_called_once()
|
||||
self.mock_session_service.append_event.assert_not_called()
|
||||
|
||||
def test_events_compaction_config_accepts_token_fields(self):
|
||||
config = EventsCompactionConfig(
|
||||
compaction_interval=2,
|
||||
overlap_size=1,
|
||||
token_threshold=50_000,
|
||||
event_retention_size=5,
|
||||
)
|
||||
self.assertEqual(config.token_threshold, 50_000)
|
||||
self.assertEqual(config.event_retention_size, 5)
|
||||
|
||||
def test_events_compaction_config_rejects_partial_token_fields(
|
||||
self,
|
||||
):
|
||||
with pytest.raises(ValidationError):
|
||||
EventsCompactionConfig(
|
||||
compaction_interval=2,
|
||||
overlap_size=1,
|
||||
token_threshold=50_000,
|
||||
)
|
||||
|
||||
def test_latest_prompt_token_count_fallback_applies_compaction(self):
|
||||
events = [
|
||||
self._create_event(1.0, 'inv1', 'a' * 40),
|
||||
self._create_event(2.0, 'inv2', 'b' * 40),
|
||||
self._create_compacted_event(1.0, 2.0, 'S'),
|
||||
self._create_event(3.0, 'inv3', 'c' * 20),
|
||||
]
|
||||
|
||||
estimated_token_count = compaction_module._latest_prompt_token_count(events)
|
||||
|
||||
# Visible text after compaction is: 'S' + ('c' * 20) = 21 chars.
|
||||
self.assertEqual(estimated_token_count, 21 // 4)
|
||||
|
||||
async def test_run_compaction_for_token_threshold_keeps_retention_events(
|
||||
self,
|
||||
):
|
||||
app = App(
|
||||
name='test',
|
||||
root_agent=Mock(spec=BaseAgent),
|
||||
events_compaction_config=EventsCompactionConfig(
|
||||
summarizer=self.mock_compactor,
|
||||
compaction_interval=999,
|
||||
overlap_size=0,
|
||||
token_threshold=50,
|
||||
event_retention_size=2,
|
||||
),
|
||||
)
|
||||
session = Session(
|
||||
app_name='test',
|
||||
user_id='u1',
|
||||
id='s1',
|
||||
events=[
|
||||
self._create_event(1.0, 'inv1', 'e1'),
|
||||
self._create_event(2.0, 'inv2', 'e2'),
|
||||
self._create_event(3.0, 'inv3', 'e3'),
|
||||
self._create_event(4.0, 'inv4', 'e4'),
|
||||
self._create_event(5.0, 'inv5', 'e5', prompt_token_count=100),
|
||||
],
|
||||
)
|
||||
|
||||
mock_compacted_event = self._create_compacted_event(
|
||||
1.0, 3.0, 'Summary inv1-inv3'
|
||||
)
|
||||
self.mock_compactor.maybe_summarize_events.return_value = (
|
||||
mock_compacted_event
|
||||
)
|
||||
|
||||
await _run_compaction_for_sliding_window(
|
||||
app, session, self.mock_session_service
|
||||
)
|
||||
|
||||
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
||||
1
|
||||
]['events']
|
||||
self.assertEqual(
|
||||
[e.invocation_id for e in compacted_events_arg],
|
||||
['inv1', 'inv2', 'inv3'],
|
||||
)
|
||||
self.mock_session_service.append_event.assert_called_once_with(
|
||||
session=session, event=mock_compacted_event
|
||||
)
|
||||
|
||||
async def test_run_compaction_for_token_threshold_seeds_previous_compaction(
|
||||
self,
|
||||
):
|
||||
app = App(
|
||||
name='test',
|
||||
root_agent=Mock(spec=BaseAgent),
|
||||
events_compaction_config=EventsCompactionConfig(
|
||||
summarizer=self.mock_compactor,
|
||||
compaction_interval=999,
|
||||
overlap_size=0,
|
||||
token_threshold=50,
|
||||
event_retention_size=2,
|
||||
),
|
||||
)
|
||||
session = Session(
|
||||
app_name='test',
|
||||
user_id='u1',
|
||||
id='s1',
|
||||
events=[
|
||||
self._create_event(1.0, 'inv1', 'e1'),
|
||||
self._create_event(2.0, 'inv2', 'e2'),
|
||||
self._create_compacted_event(1.0, 2.0, 'Summary 1-2'),
|
||||
self._create_event(3.0, 'inv3', 'e3'),
|
||||
self._create_event(4.0, 'inv4', 'e4'),
|
||||
self._create_event(5.0, 'inv5', 'e5'),
|
||||
self._create_event(6.0, 'inv6', 'e6', prompt_token_count=100),
|
||||
],
|
||||
)
|
||||
|
||||
mock_compacted_event = self._create_compacted_event(1.0, 4.0, 'Summary 1-4')
|
||||
self.mock_compactor.maybe_summarize_events.return_value = (
|
||||
mock_compacted_event
|
||||
)
|
||||
|
||||
await _run_compaction_for_sliding_window(
|
||||
app, session, self.mock_session_service
|
||||
)
|
||||
|
||||
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
||||
1
|
||||
]['events']
|
||||
self.assertEqual(
|
||||
[e.content.parts[0].text for e in compacted_events_arg],
|
||||
['Summary 1-2', 'e3', 'e4'],
|
||||
)
|
||||
self.assertEqual(compacted_events_arg[0].timestamp, 1.0)
|
||||
self.assertEqual(
|
||||
[e.invocation_id for e in compacted_events_arg[1:]],
|
||||
['inv3', 'inv4'],
|
||||
)
|
||||
self.mock_session_service.append_event.assert_called_once_with(
|
||||
session=session, event=mock_compacted_event
|
||||
)
|
||||
|
||||
async def test_run_compaction_for_token_threshold_with_zero_retention(
|
||||
self,
|
||||
):
|
||||
app = App(
|
||||
name='test',
|
||||
root_agent=Mock(spec=BaseAgent),
|
||||
events_compaction_config=EventsCompactionConfig(
|
||||
summarizer=self.mock_compactor,
|
||||
compaction_interval=999,
|
||||
overlap_size=0,
|
||||
token_threshold=50,
|
||||
event_retention_size=0,
|
||||
),
|
||||
)
|
||||
session = Session(
|
||||
app_name='test',
|
||||
user_id='u1',
|
||||
id='s1',
|
||||
events=[
|
||||
self._create_event(1.0, 'inv1', 'e1'),
|
||||
self._create_event(2.0, 'inv2', 'e2'),
|
||||
self._create_event(3.0, 'inv3', 'e3', prompt_token_count=100),
|
||||
],
|
||||
)
|
||||
|
||||
mock_compacted_event = self._create_compacted_event(
|
||||
1.0, 3.0, 'Summary inv1-inv3'
|
||||
)
|
||||
self.mock_compactor.maybe_summarize_events.return_value = (
|
||||
mock_compacted_event
|
||||
)
|
||||
|
||||
await _run_compaction_for_sliding_window(
|
||||
app, session, self.mock_session_service
|
||||
)
|
||||
|
||||
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
||||
1
|
||||
]['events']
|
||||
self.assertEqual(
|
||||
[e.invocation_id for e in compacted_events_arg],
|
||||
['inv1', 'inv2', 'inv3'],
|
||||
)
|
||||
self.mock_session_service.append_event.assert_called_once_with(
|
||||
session=session, event=mock_compacted_event
|
||||
)
|
||||
|
||||
async def test_run_compaction_for_token_threshold_with_retention_and_overlap(
|
||||
self,
|
||||
):
|
||||
app = App(
|
||||
name='test',
|
||||
root_agent=Mock(spec=BaseAgent),
|
||||
events_compaction_config=EventsCompactionConfig(
|
||||
summarizer=self.mock_compactor,
|
||||
compaction_interval=999,
|
||||
overlap_size=0,
|
||||
token_threshold=50,
|
||||
event_retention_size=3,
|
||||
),
|
||||
)
|
||||
session = Session(
|
||||
app_name='test',
|
||||
user_id='u1',
|
||||
id='s1',
|
||||
events=[
|
||||
self._create_event(1.0, 'inv1', 'e1'),
|
||||
self._create_event(2.0, 'inv2', 'e2'),
|
||||
self._create_event(3.0, 'inv3', 'e3'),
|
||||
self._create_event(4.0, 'inv4', 'e4'),
|
||||
self._create_compacted_event(
|
||||
1.0, 1.0, 'Summary 1', appended_ts=5.0
|
||||
),
|
||||
self._create_event(6.0, 'inv6', 'e6'),
|
||||
self._create_event(7.0, 'inv7', 'e7'),
|
||||
self._create_compacted_event(
|
||||
1.0, 3.0, 'Summary 1-3', appended_ts=8.0
|
||||
),
|
||||
self._create_event(9.0, 'inv9', 'e9', prompt_token_count=100),
|
||||
],
|
||||
)
|
||||
|
||||
mock_compacted_event = self._create_compacted_event(1.0, 4.0, 'Summary 1-4')
|
||||
self.mock_compactor.maybe_summarize_events.return_value = (
|
||||
mock_compacted_event
|
||||
)
|
||||
|
||||
await _run_compaction_for_sliding_window(
|
||||
app, session, self.mock_session_service
|
||||
)
|
||||
|
||||
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
||||
1
|
||||
]['events']
|
||||
self.assertEqual(
|
||||
[e.content.parts[0].text for e in compacted_events_arg],
|
||||
['Summary 1-3', 'e4'],
|
||||
)
|
||||
self.assertEqual(compacted_events_arg[0].timestamp, 1.0)
|
||||
self.assertEqual(compacted_events_arg[1].invocation_id, 'inv4')
|
||||
self.mock_session_service.append_event.assert_called_once_with(
|
||||
session=session, event=mock_compacted_event
|
||||
)
|
||||
|
||||
def test_get_contents_with_multiple_compactions(self):
|
||||
|
||||
# Event timestamps: 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0
|
||||
@@ -262,6 +519,45 @@ class TestCompaction(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(actual_texts, expected_texts)
|
||||
# Verify timestamps are in order
|
||||
|
||||
def test_get_contents_subsumed_compaction_is_hidden(self):
|
||||
events = [
|
||||
self._create_event(1.0, 'inv1', 'Event 1'),
|
||||
self._create_event(2.0, 'inv2', 'Event 2'),
|
||||
self._create_event(3.0, 'inv3', 'Event 3'),
|
||||
self._create_event(4.0, 'inv4', 'Event 4'),
|
||||
self._create_compacted_event(1.0, 1.0, 'Summary 1'),
|
||||
self._create_event(6.0, 'inv6', 'Event 6'),
|
||||
self._create_event(7.0, 'inv7', 'Event 7'),
|
||||
self._create_compacted_event(1.0, 3.0, 'Summary 1-3'),
|
||||
self._create_event(9.0, 'inv9', 'Event 9'),
|
||||
]
|
||||
|
||||
result_contents = contents._get_contents(None, events)
|
||||
expected_texts = [
|
||||
'Summary 1-3',
|
||||
'Event 4',
|
||||
'Event 6',
|
||||
'Event 7',
|
||||
'Event 9',
|
||||
]
|
||||
actual_texts = [c.parts[0].text for c in result_contents]
|
||||
self.assertEqual(actual_texts, expected_texts)
|
||||
|
||||
def test_get_contents_compaction_appended_late_keeps_newer_events(self):
|
||||
events = [
|
||||
self._create_event(1.0, 'inv1', 'Event 1'),
|
||||
self._create_event(2.0, 'inv2', 'Event 2'),
|
||||
self._create_event(3.0, 'inv3', 'Event 3'),
|
||||
self._create_event(4.0, 'inv4', 'Event 4'),
|
||||
self._create_event(5.0, 'inv5', 'Event 5'),
|
||||
self._create_compacted_event(1.0, 3.0, 'Summary 1-3', appended_ts=6.0),
|
||||
]
|
||||
|
||||
result_contents = contents._get_contents(None, events)
|
||||
expected_texts = ['Summary 1-3', 'Event 4', 'Event 5']
|
||||
actual_texts = [c.parts[0].text for c in result_contents]
|
||||
self.assertEqual(actual_texts, expected_texts)
|
||||
|
||||
def test_get_contents_no_compaction(self):
|
||||
|
||||
events = [
|
||||
|
||||
Reference in New Issue
Block a user