wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
1974 lines
67 KiB
Python
1974 lines
67 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import unittest
|
|
from unittest.mock import AsyncMock
|
|
from unittest.mock import Mock
|
|
|
|
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.base_events_summarizer import BaseEventsSummarizer
|
|
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.auth.auth_schemes import CustomAuthScheme
|
|
from google.adk.auth.auth_tool import AuthConfig
|
|
from google.adk.events.event import Event
|
|
from google.adk.events.event_actions import EventActions
|
|
from google.adk.events.event_actions import EventCompaction
|
|
from google.adk.flows.llm_flows import contents as _contents
|
|
from google.adk.sessions.base_session_service import BaseSessionService
|
|
from google.adk.sessions.session import Session
|
|
from google.adk.tools.tool_confirmation import ToolConfirmation
|
|
from google.genai import types
|
|
from google.genai.types import Content
|
|
from google.genai.types import Part
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
|
from pydantic import ValidationError
|
|
import pytest
|
|
|
|
|
|
class _StubSummarizer(BaseEventsSummarizer):
|
|
|
|
def __init__(self, compacted_event: Event | None):
|
|
self._compacted_event = compacted_event
|
|
|
|
async def maybe_summarize_events(
|
|
self, *, events: list[Event]
|
|
) -> Event | None:
|
|
del events
|
|
return self._compacted_event
|
|
|
|
|
|
def _create_trace_test_event(
|
|
*,
|
|
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_trace_compacted_event(
|
|
*, start_ts: float, end_ts: float, summary_text: str
|
|
) -> Event:
|
|
compaction = EventCompaction(
|
|
start_timestamp=start_ts,
|
|
end_timestamp=end_ts,
|
|
compacted_content=Content(role='model', parts=[Part(text=summary_text)]),
|
|
)
|
|
return Event(
|
|
id='compacted-event-id',
|
|
timestamp=end_ts,
|
|
author='compactor',
|
|
content=compaction.compacted_content,
|
|
actions=EventActions(compaction=compaction),
|
|
invocation_id='compacted-invocation-id',
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def span_exporter(monkeypatch: pytest.MonkeyPatch) -> InMemorySpanExporter:
|
|
tracer_provider = TracerProvider()
|
|
span_exporter = InMemorySpanExporter()
|
|
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
|
|
real_tracer = tracer_provider.get_tracer(__name__)
|
|
monkeypatch.setattr(
|
|
compaction_module.tracer,
|
|
'start_as_current_span',
|
|
real_tracer.start_as_current_span,
|
|
)
|
|
return span_exporter
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'env_variables', ['GOOGLE_AI', 'VERTEX'], indirect=True
|
|
)
|
|
class TestCompaction(unittest.IsolatedAsyncioTestCase):
|
|
|
|
def setUp(self):
|
|
self.mock_session_service = AsyncMock(spec=BaseSessionService)
|
|
self.mock_compactor = AsyncMock(spec=LlmEventSummarizer)
|
|
|
|
def _create_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
text: str,
|
|
prompt_token_count: int | None = None,
|
|
thought: bool = False,
|
|
) -> 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, thought=thought)]),
|
|
usage_metadata=usage_metadata,
|
|
)
|
|
|
|
def _create_function_call_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
function_call_id: str,
|
|
) -> Event:
|
|
return Event(
|
|
timestamp=timestamp,
|
|
invocation_id=invocation_id,
|
|
author='agent',
|
|
content=Content(
|
|
role='model',
|
|
parts=[
|
|
Part(
|
|
function_call=types.FunctionCall(
|
|
id=function_call_id, name='tool', args={}
|
|
)
|
|
)
|
|
],
|
|
),
|
|
)
|
|
|
|
def _create_function_response_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
function_call_id: 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='agent',
|
|
content=Content(
|
|
role='user',
|
|
parts=[
|
|
Part(
|
|
function_response=types.FunctionResponse(
|
|
id=function_call_id,
|
|
name='tool',
|
|
response={'result': 'ok'},
|
|
)
|
|
)
|
|
],
|
|
),
|
|
usage_metadata=usage_metadata,
|
|
)
|
|
|
|
def _create_compacted_event(
|
|
self,
|
|
start_ts: float,
|
|
end_ts: float,
|
|
summary_text: str,
|
|
appended_ts: float | None = None,
|
|
) -> Event:
|
|
compaction = EventCompaction(
|
|
start_timestamp=start_ts,
|
|
end_timestamp=end_ts,
|
|
compacted_content=Content(
|
|
role='model', parts=[Part(text=summary_text)]
|
|
),
|
|
)
|
|
return Event(
|
|
timestamp=appended_ts if appended_ts is not None else end_ts,
|
|
author='compactor',
|
|
content=compaction.compacted_content,
|
|
actions=EventActions(compaction=compaction),
|
|
invocation_id=Event.new_id(),
|
|
)
|
|
|
|
async def _run_sliding_window(self, app, session, session_service, **kwargs):
|
|
"""Drains the sliding-window generator, appending like the runner loop.
|
|
|
|
Compaction now yields its event instead of appending it, so the runner is
|
|
the single append site. This mirrors that behavior so the existing append
|
|
assertions still exercise the produce-then-persist path.
|
|
"""
|
|
async for compaction_event in _run_compaction_for_sliding_window(
|
|
app, session, session_service, **kwargs
|
|
):
|
|
await session_service.append_event(
|
|
session=session, event=compaction_event
|
|
)
|
|
|
|
async def test_run_compaction_for_sliding_window_no_events(self):
|
|
app = App(name='test', root_agent=Mock(spec=BaseAgent))
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=[])
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
self.mock_compactor.maybe_summarize_events.assert_not_called()
|
|
self.mock_session_service.append_event.assert_not_called()
|
|
|
|
async def test_sliding_window_yields_event_without_appending(self):
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=1,
|
|
),
|
|
)
|
|
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'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 4.0, 'Summary inv1-inv4'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
yielded = [
|
|
event
|
|
async for event in _run_compaction_for_sliding_window(
|
|
app, session, self.mock_session_service
|
|
)
|
|
]
|
|
|
|
# The compaction event is yielded to the caller (the runner loop), and the
|
|
# function itself never appends -- persistence is the runner's job.
|
|
self.assertEqual(yielded, [mock_compacted_event])
|
|
self.mock_session_service.append_event.assert_not_called()
|
|
|
|
async def test_sliding_window_yields_nothing_when_no_compaction(self):
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=3,
|
|
overlap_size=1,
|
|
),
|
|
)
|
|
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'),
|
|
],
|
|
)
|
|
|
|
yielded = [
|
|
event
|
|
async for event in _run_compaction_for_sliding_window(
|
|
app, session, self.mock_session_service
|
|
)
|
|
]
|
|
|
|
self.assertEqual(yielded, [])
|
|
self.mock_session_service.append_event.assert_not_called()
|
|
|
|
async def test_run_compaction_for_sliding_window_not_enough_new_invocations(
|
|
self,
|
|
):
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=3,
|
|
overlap_size=1,
|
|
),
|
|
)
|
|
# Only two new invocations ('inv1', 'inv2'), less than compaction_interval=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'),
|
|
],
|
|
)
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
self.mock_compactor.maybe_summarize_events.assert_not_called()
|
|
self.mock_session_service.append_event.assert_not_called()
|
|
|
|
async def test_run_compaction_for_sliding_window_first_compaction(self):
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=1,
|
|
),
|
|
)
|
|
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'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 4.0, 'Summary inv1-inv4'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
# Expected events to compact: inv1, inv2, inv3, inv4
|
|
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', 'inv4'],
|
|
)
|
|
self.mock_session_service.append_event.assert_called_once_with(
|
|
session=session, event=mock_compacted_event
|
|
)
|
|
|
|
async def test_run_compaction_for_sliding_window_with_overlap(self):
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=1,
|
|
),
|
|
)
|
|
# inv1-inv2 are already compacted. Last compacted end timestamp is 2.0.
|
|
initial_events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_event(2.0, 'inv2', 'e2'),
|
|
self._create_compacted_event(1.0, 2.0, 'Summary inv1-inv2'),
|
|
]
|
|
# Add new invocations inv3, inv4, inv5
|
|
new_events = [
|
|
self._create_event(3.0, 'inv3', 'e3'),
|
|
self._create_event(4.0, 'inv4', 'e4'),
|
|
self._create_event(5.0, 'inv5', 'e5'),
|
|
]
|
|
session = Session(
|
|
app_name='test',
|
|
user_id='u1',
|
|
id='s1',
|
|
events=initial_events + new_events,
|
|
)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
2.0, 5.0, 'Summary inv2-inv5'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
# New invocations are inv3, inv4, inv5 (3 new) > threshold (2).
|
|
# Overlap size is 1, so start from 1 inv before inv3, which is inv2.
|
|
# Compact range: inv2 to inv5.
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
self.assertEqual(
|
|
[e.invocation_id for e in compacted_events_arg],
|
|
['inv2', 'inv3', 'inv4', 'inv5'],
|
|
)
|
|
self.mock_session_service.append_event.assert_called_once_with(
|
|
session=session, event=mock_compacted_event
|
|
)
|
|
|
|
async def test_run_compaction_for_sliding_window_no_compaction_event_returned(
|
|
self,
|
|
):
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=1,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
events = [self._create_event(1.0, 'inv1', 'e1')]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
self.mock_compactor.maybe_summarize_events.return_value = None
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
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.compaction_interval, 2)
|
|
self.assertEqual(config.overlap_size, 1)
|
|
self.assertEqual(config.token_threshold, 50_000)
|
|
self.assertEqual(config.event_retention_size, 5)
|
|
|
|
def test_events_compaction_config_accepts_sliding_window_fields(self):
|
|
config = EventsCompactionConfig(
|
|
compaction_interval=2,
|
|
overlap_size=1,
|
|
)
|
|
self.assertEqual(config.compaction_interval, 2)
|
|
self.assertEqual(config.overlap_size, 1)
|
|
self.assertIsNone(config.token_threshold)
|
|
self.assertIsNone(config.event_retention_size)
|
|
|
|
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_events_compaction_config_rejects_partial_sliding_fields(
|
|
self,
|
|
):
|
|
with pytest.raises(ValidationError):
|
|
EventsCompactionConfig(
|
|
compaction_interval=2,
|
|
)
|
|
|
|
with pytest.raises(ValidationError):
|
|
EventsCompactionConfig(
|
|
overlap_size=0,
|
|
)
|
|
|
|
def test_events_compaction_config_rejects_missing_modes(self):
|
|
with pytest.raises(ValidationError):
|
|
EventsCompactionConfig()
|
|
|
|
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)
|
|
|
|
def test_latest_prompt_token_count_fallback_uses_effective_contents(self):
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'visible'),
|
|
Event(
|
|
timestamp=2.0,
|
|
invocation_id='inv2',
|
|
author='model',
|
|
content=Content(
|
|
role='model',
|
|
parts=[Part(text='hidden-thought', thought=True)],
|
|
),
|
|
),
|
|
]
|
|
|
|
estimated_token_count = compaction_module._latest_prompt_token_count(events)
|
|
|
|
# Thought-only events are filtered by contents processing.
|
|
self.assertEqual(estimated_token_count, len('visible') // 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 self._run_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_keeps_tool_call_pair(
|
|
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=1,
|
|
),
|
|
)
|
|
session = Session(
|
|
app_name='test',
|
|
user_id='u1',
|
|
id='s1',
|
|
events=[
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'tool-call-1'),
|
|
self._create_function_response_event(
|
|
3.0,
|
|
'inv2',
|
|
'tool-call-1',
|
|
prompt_token_count=100,
|
|
),
|
|
],
|
|
)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 1.0, 'Summary inv1'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_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'],
|
|
)
|
|
self.mock_session_service.append_event.assert_called_once_with(
|
|
session=session, event=mock_compacted_event
|
|
)
|
|
|
|
async def test_run_compaction_for_token_threshold_equal_threshold_compacts(
|
|
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=100,
|
|
event_retention_size=1,
|
|
),
|
|
)
|
|
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', prompt_token_count=100),
|
|
],
|
|
)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 1.0, 'Summary inv1'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_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'],
|
|
)
|
|
self.mock_session_service.append_event.assert_called_once_with(
|
|
session=session, event=mock_compacted_event
|
|
)
|
|
|
|
async def test_run_compaction_skip_token_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=1,
|
|
),
|
|
)
|
|
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', prompt_token_count=100),
|
|
],
|
|
)
|
|
|
|
await self._run_sliding_window(
|
|
app,
|
|
session,
|
|
self.mock_session_service,
|
|
skip_token_compaction=True,
|
|
)
|
|
|
|
self.mock_compactor.maybe_summarize_events.assert_not_called()
|
|
self.mock_session_service.append_event.assert_not_called()
|
|
|
|
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 self._run_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 self._run_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 self._run_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
|
|
)
|
|
|
|
async def test_run_compaction_for_token_threshold_uses_latest_ordered_seed(
|
|
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=1,
|
|
),
|
|
)
|
|
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'),
|
|
self._create_event(15.0, 'inv6', 'e6'),
|
|
self._create_event(20.0, 'inv7', 'e7'),
|
|
self._create_compacted_event(
|
|
15.0, 20.0, 'Summary 15-20', appended_ts=21.0
|
|
),
|
|
self._create_compacted_event(
|
|
1.0, 5.0, 'Summary 1-5', appended_ts=22.0
|
|
),
|
|
self._create_event(23.0, 'inv8', 'e8'),
|
|
self._create_event(24.0, 'inv9', 'e9', prompt_token_count=120),
|
|
],
|
|
)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 23.0, 'Summary 1-23'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
self.assertEqual(
|
|
compacted_events_arg[0].content.parts[0].text, 'Summary 1-5'
|
|
)
|
|
self.assertEqual(
|
|
[e.invocation_id for e in compacted_events_arg[1:]],
|
|
['inv6', 'inv7', 'inv8'],
|
|
)
|
|
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
|
|
# Compaction 1: covers 1.0 to 4.0 (summary at 4.0)
|
|
# Compaction 2: covers 6.0 to 9.0 (summary at 9.0)
|
|
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, 4.0, 'Summary 1-4'),
|
|
self._create_event(5.0, 'inv5', 'Event 5'),
|
|
self._create_event(6.0, 'inv6', 'Event 6'),
|
|
self._create_event(7.0, 'inv7', 'Event 7'),
|
|
self._create_event(8.0, 'inv8', 'Event 8'),
|
|
self._create_event(9.0, 'inv9', 'Event 9'),
|
|
self._create_compacted_event(6.0, 9.0, 'Summary 6-9'),
|
|
self._create_event(10.0, 'inv10', 'Event 10'),
|
|
]
|
|
|
|
result_contents = _contents._get_contents(None, events)
|
|
|
|
# Expected contents:
|
|
# Summary 1-4 (at timestamp 4.0)
|
|
# Event 5 (at timestamp 5.0)
|
|
# Summary 6-9 (at timestamp 9.0)
|
|
# Event 10 (at timestamp 10.0)
|
|
expected_texts = [
|
|
'Summary 1-4',
|
|
'Event 5',
|
|
'Summary 6-9',
|
|
'Event 10',
|
|
]
|
|
actual_texts = [c.parts[0].text for c in result_contents]
|
|
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 = [
|
|
self._create_event(1.0, 'inv1', 'Event 1'),
|
|
self._create_event(2.0, 'inv2', 'Event 2'),
|
|
self._create_event(3.0, 'inv3', 'Event 3'),
|
|
]
|
|
|
|
result_contents = _contents._get_contents(None, events)
|
|
expected_texts = ['Event 1', 'Event 2', 'Event 3']
|
|
actual_texts = [c.parts[0].text for c in result_contents]
|
|
self.assertEqual(actual_texts, expected_texts)
|
|
|
|
def test_get_contents_single_compaction_at_start(self):
|
|
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'Event 1'),
|
|
self._create_event(2.0, 'inv2', 'Event 2'),
|
|
self._create_compacted_event(1.0, 2.0, 'Summary 1-2'),
|
|
self._create_event(3.0, 'inv3', 'Event 3'),
|
|
]
|
|
|
|
result_contents = _contents._get_contents(None, events)
|
|
expected_texts = ['Summary 1-2', 'Event 3']
|
|
actual_texts = [c.parts[0].text for c in result_contents]
|
|
self.assertEqual(actual_texts, expected_texts)
|
|
|
|
def test_get_contents_single_compaction_in_middle(self):
|
|
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'Event 1'),
|
|
self._create_event(2.0, 'inv2', 'Event 2'),
|
|
self._create_compacted_event(1.0, 2.0, 'Summary 1-2'),
|
|
self._create_event(3.0, 'inv3', 'Event 3'),
|
|
self._create_event(4.0, 'inv4', 'Event 4'),
|
|
self._create_compacted_event(3.0, 4.0, 'Summary 3-4'),
|
|
self._create_event(5.0, 'inv5', 'Event 5'),
|
|
]
|
|
|
|
result_contents = _contents._get_contents(None, events)
|
|
expected_texts = ['Summary 1-2', 'Summary 3-4', 'Event 5']
|
|
actual_texts = [c.parts[0].text for c in result_contents]
|
|
self.assertEqual(actual_texts, expected_texts)
|
|
|
|
def test_get_contents_compaction_at_end(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_compacted_event(2.0, 3.0, 'Summary 2-3'),
|
|
]
|
|
|
|
result_contents = _contents._get_contents(None, events)
|
|
expected_texts = ['Event 1', 'Summary 2-3']
|
|
actual_texts = [c.parts[0].text for c in result_contents]
|
|
self.assertEqual(actual_texts, expected_texts)
|
|
|
|
def test_get_contents_compaction_at_beginning(self):
|
|
|
|
events = [
|
|
self._create_compacted_event(1.0, 2.0, 'Summary 1-2'),
|
|
self._create_event(3.0, 'inv3', 'Event 3'),
|
|
self._create_event(4.0, 'inv4', 'Event 4'),
|
|
]
|
|
|
|
result_contents = _contents._get_contents(None, events)
|
|
expected_texts = ['Summary 1-2', 'Event 3', 'Event 4']
|
|
actual_texts = [c.parts[0].text for c in result_contents]
|
|
self.assertEqual(actual_texts, expected_texts)
|
|
|
|
async def test_sliding_window_excludes_pending_function_call_events(self):
|
|
"""Sliding-window compaction stops before pending function calls."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# inv1: normal text, inv2: pending function call (no response)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'pending-call-1'),
|
|
self._create_event(3.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 3.0, 'Summary without pending'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_sliding_window_pending_function_call_remains_in_contents(
|
|
self,
|
|
):
|
|
"""Sliding-window compaction keeps pending tool calls visible in history."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'pending-call-1'),
|
|
self._create_event(3.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
self.mock_compactor.maybe_summarize_events.side_effect = (
|
|
lambda *, events: self._create_compacted_event(
|
|
events[0].timestamp,
|
|
events[-1].timestamp,
|
|
'Summary safe prefix',
|
|
)
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
appended_event = self.mock_session_service.append_event.call_args[1][
|
|
'event'
|
|
]
|
|
self.assertEqual(appended_event.actions.compaction.start_timestamp, 1.0)
|
|
self.assertEqual(appended_event.actions.compaction.end_timestamp, 1.0)
|
|
|
|
result_contents = _contents._get_contents(None, events + [appended_event])
|
|
self.assertEqual(result_contents[0].parts[0].text, 'Summary safe prefix')
|
|
self.assertEqual(
|
|
result_contents[1].parts[0].function_call.name,
|
|
'tool',
|
|
)
|
|
self.assertEqual(result_contents[2].parts[0].text, 'e3')
|
|
|
|
async def test_token_threshold_excludes_pending_function_call_events(self):
|
|
"""Token-threshold compaction stays contiguous before pending calls."""
|
|
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,
|
|
),
|
|
)
|
|
# inv1: text, inv2: pending function call, inv3: text with token count
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'pending-call-1'),
|
|
self._create_event(3.0, 'inv3', 'e3', prompt_token_count=100),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 1.0, 'Summary inv1'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_token_threshold_pending_function_call_remains_in_contents(
|
|
self,
|
|
):
|
|
"""Token-threshold compaction keeps pending tool calls visible."""
|
|
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,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'pending-call-1'),
|
|
self._create_event(3.0, 'inv3', 'e3', prompt_token_count=100),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
self.mock_compactor.maybe_summarize_events.side_effect = (
|
|
lambda *, events: self._create_compacted_event(
|
|
events[0].timestamp,
|
|
events[-1].timestamp,
|
|
'Summary safe prefix',
|
|
)
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
appended_event = self.mock_session_service.append_event.call_args[1][
|
|
'event'
|
|
]
|
|
self.assertEqual(appended_event.actions.compaction.start_timestamp, 1.0)
|
|
self.assertEqual(appended_event.actions.compaction.end_timestamp, 1.0)
|
|
|
|
result_contents = _contents._get_contents(None, events + [appended_event])
|
|
self.assertEqual(result_contents[0].parts[0].text, 'Summary safe prefix')
|
|
self.assertEqual(
|
|
result_contents[1].parts[0].function_call.name,
|
|
'tool',
|
|
)
|
|
self.assertEqual(result_contents[2].parts[0].text, 'e3')
|
|
|
|
async def test_completed_function_call_pair_is_still_compacted(self):
|
|
"""Completed function call/response pairs must still be compacted."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# inv1: text, inv2: completed call+response pair, inv3: text
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'completed-call-1'),
|
|
self._create_function_response_event(3.0, 'inv2', 'completed-call-1'),
|
|
self._create_event(4.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 4.0, 'Summary with completed pair'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# Both the call and response events for inv2 should be compacted.
|
|
self.assertIn('inv1', compacted_inv_ids)
|
|
self.assertEqual(compacted_inv_ids.count('inv2'), 2)
|
|
self.assertIn('inv3', compacted_inv_ids)
|
|
|
|
def _create_hitl_confirmation_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
function_call_id: str,
|
|
) -> Event:
|
|
"""Creates a function response event with a tool confirmation request."""
|
|
return Event(
|
|
timestamp=timestamp,
|
|
invocation_id=invocation_id,
|
|
author='agent',
|
|
content=Content(
|
|
role='user',
|
|
parts=[
|
|
Part(
|
|
function_response=types.FunctionResponse(
|
|
id=function_call_id,
|
|
name='tool',
|
|
response={
|
|
'error': 'This tool call requires confirmation.'
|
|
},
|
|
)
|
|
)
|
|
],
|
|
),
|
|
actions=EventActions(
|
|
requested_tool_confirmations={
|
|
function_call_id: ToolConfirmation(
|
|
hint='Please confirm this action.'
|
|
)
|
|
},
|
|
),
|
|
)
|
|
|
|
def _create_hitl_auth_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
function_call_id: str,
|
|
) -> Event:
|
|
"""Creates a function response event with an auth credential request."""
|
|
return Event(
|
|
timestamp=timestamp,
|
|
invocation_id=invocation_id,
|
|
author='agent',
|
|
content=Content(
|
|
role='user',
|
|
parts=[
|
|
Part(
|
|
function_response=types.FunctionResponse(
|
|
id=function_call_id,
|
|
name='tool',
|
|
response={'error': 'Auth required.'},
|
|
)
|
|
)
|
|
],
|
|
),
|
|
actions=EventActions(
|
|
requested_auth_configs={
|
|
function_call_id: AuthConfig(
|
|
auth_scheme=CustomAuthScheme(type='custom'),
|
|
)
|
|
},
|
|
),
|
|
)
|
|
|
|
async def test_sliding_window_excludes_hitl_confirmation_events(self):
|
|
"""Sliding-window compaction stops before tool confirmation events."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# inv1: text, inv2: call + HITL confirmation response, inv3: text
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_hitl_confirmation_event(3.0, 'inv2', 'call-1'),
|
|
self._create_event(4.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 2.0, 'Summary before hitl'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# inv2's tool call is still awaiting the final response,
|
|
# so compaction won't summarize it; only the
|
|
# already settled inv1 is compacted.
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_sliding_window_excludes_hitl_auth_events(self):
|
|
"""Sliding-window compaction stops before auth credential events."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_hitl_auth_event(3.0, 'inv2', 'call-1'),
|
|
self._create_event(4.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 2.0, 'Summary before auth'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# inv2's tool call is still awaiting auth approval -- an unfinished
|
|
# call/response pair -- so compaction won't summarize it; only the
|
|
# already settled inv1 is compacted.
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_token_threshold_excludes_hitl_confirmation_events(self):
|
|
"""Token-threshold compaction stops before tool confirmation events."""
|
|
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,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_hitl_confirmation_event(3.0, 'inv2', 'call-1'),
|
|
self._create_event(4.0, 'inv3', 'e3', prompt_token_count=100),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 2.0, 'Summary inv1-inv2'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# inv2's tool call is still awaiting confirmation -- an unfinished
|
|
# call/response pair -- so compaction won't summarize it; only the
|
|
# already settled inv1 is compacted.
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_token_threshold_excludes_hitl_auth_events(self):
|
|
"""Token-threshold compaction stops before auth credential events."""
|
|
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,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_hitl_auth_event(3.0, 'inv2', 'call-1'),
|
|
self._create_event(4.0, 'inv3', 'e3', prompt_token_count=100),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 2.0, 'Summary inv1-inv2'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# inv2's tool call is still awaiting auth approval -- an unfinished
|
|
# call/response pair -- so compaction won't summarize it; only the
|
|
# already settled inv1 is compacted.
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_hitl_event_at_start_blocks_all_compaction(self):
|
|
"""If the first candidate event has HITL, nothing is compacted."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# The very first event is an HITL confirmation (no preceding function call).
|
|
events = [
|
|
self._create_hitl_confirmation_event(1.0, 'inv1', 'call-1'),
|
|
self._create_event(2.0, 'inv2', 'e2'),
|
|
self._create_event(3.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
self.mock_compactor.maybe_summarize_events.assert_not_called()
|
|
self.mock_session_service.append_event.assert_not_called()
|
|
|
|
async def test_events_before_hitl_are_still_compacted(self):
|
|
"""Events before the HITL event are compacted normally."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# inv1, inv2: text events, inv3: call + HITL confirmation, inv4: text
|
|
# The HITL event at index 3 blocks compaction; events before it are safe.
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_event(2.0, 'inv2', 'e2'),
|
|
self._create_function_call_event(3.0, 'inv3', 'call-1'),
|
|
self._create_hitl_confirmation_event(4.0, 'inv3', 'call-1'),
|
|
self._create_event(5.0, 'inv4', 'e4'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
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 self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# inv3's tool call is still awaiting confirmation -- an unfinished
|
|
# call/response pair -- so compaction stops before it; the settled inv1
|
|
# and inv2 are compacted.
|
|
self.assertEqual(compacted_inv_ids, ['inv1', 'inv2'])
|
|
|
|
async def test_resolved_hitl_confirmation_is_compactable(self):
|
|
"""A HITL confirmation followed by a resolved tool response is compactable."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# inv1: text, inv2: call + HITL request + resolved response (same call-1
|
|
# id), inv3: text. The resolved HITL is safe to compact.
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_hitl_confirmation_event(3.0, 'inv2', 'call-1'),
|
|
self._create_function_response_event(4.0, 'inv2', 'call-1'),
|
|
self._create_event(5.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 5.0, 'Summary including resolved hitl'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# Resolved HITL doesn't block; all events through inv3 compact together.
|
|
self.assertEqual(
|
|
compacted_inv_ids, ['inv1', 'inv2', 'inv2', 'inv2', 'inv3']
|
|
)
|
|
|
|
async def test_resolved_hitl_auth_is_compactable(self):
|
|
"""A HITL auth request followed by a resolved tool response is compactable."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_hitl_auth_event(3.0, 'inv2', 'call-1'),
|
|
self._create_function_response_event(4.0, 'inv2', 'call-1'),
|
|
self._create_event(5.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 5.0, 'Summary including resolved auth'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
self.assertEqual(
|
|
compacted_inv_ids, ['inv1', 'inv2', 'inv2', 'inv2', 'inv3']
|
|
)
|
|
|
|
def _create_request_confirmation_call_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
request_confirmation_id: str,
|
|
original_function_call_id: str,
|
|
) -> Event:
|
|
"""Creates the synthetic adk_request_confirmation function-call event."""
|
|
# Mirrors functions.generate_request_confirmation_event: real ADK emits a
|
|
# separate event whose function call has its own distinct id (registered in
|
|
# long_running_tool_ids) and only references the original call id in its
|
|
# args. See tests/unittests/runners/test_run_tool_confirmation.py.
|
|
return Event(
|
|
timestamp=timestamp,
|
|
invocation_id=invocation_id,
|
|
author='agent',
|
|
content=Content(
|
|
role='model',
|
|
parts=[
|
|
Part(
|
|
function_call=types.FunctionCall(
|
|
id=request_confirmation_id,
|
|
name='adk_request_confirmation',
|
|
args={
|
|
'originalFunctionCall': {
|
|
'id': original_function_call_id
|
|
}
|
|
},
|
|
)
|
|
)
|
|
],
|
|
),
|
|
long_running_tool_ids={request_confirmation_id},
|
|
)
|
|
|
|
def _create_request_confirmation_response_event(
|
|
self,
|
|
timestamp: float,
|
|
invocation_id: str,
|
|
request_confirmation_id: str,
|
|
) -> Event:
|
|
"""Creates the function_response that resolves the confirmation call."""
|
|
return Event(
|
|
timestamp=timestamp,
|
|
invocation_id=invocation_id,
|
|
author='user',
|
|
content=Content(
|
|
role='user',
|
|
parts=[
|
|
Part(
|
|
function_response=types.FunctionResponse(
|
|
id=request_confirmation_id,
|
|
name='adk_request_confirmation',
|
|
response={'confirmed': True},
|
|
)
|
|
)
|
|
],
|
|
),
|
|
)
|
|
|
|
async def test_sliding_window_real_hitl_shape_blocks_compaction(self):
|
|
"""Faithful 3-event HITL turn (two ids) that blocks compaction.
|
|
|
|
Unlike the other HITL tests, this mirrors the event stream emitted by the
|
|
ADK runtime in functions.generate_request_confirmation_event: a
|
|
confirmation-required call produces function_call(call-1),
|
|
function_call(adk_request_confirmation) with its own distinct
|
|
client-generated id (registered in long_running_tool_ids), then the
|
|
placeholder function_response(call-1) requesting confirmation. Both the tool
|
|
call and the still-unanswered confirmation call are open, so nothing past
|
|
inv1 may be compacted.
|
|
"""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_request_confirmation_call_event(
|
|
3.0, 'inv2', 'confirm-1', 'call-1'
|
|
),
|
|
self._create_hitl_confirmation_event(4.0, 'inv2', 'call-1'),
|
|
self._create_event(5.0, 'inv3', 'e3'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 1.0, 'Summary inv1'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# Only inv1 is self-contained: call-1 awaits confirmation and the
|
|
# adk_request_confirmation call (confirm-1) has no response in the window.
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_sliding_window_real_hitl_shape_resolved_is_compactable(self):
|
|
"""Faithful resolved HITL turn (both ids closed) compacts fully.
|
|
|
|
Adds the two resolving responses seen on resume:
|
|
function_response(adk_request_confirmation) closes the confirmation call and
|
|
function_response(call-1) returns the tool result. With every obligation
|
|
closed, the whole span is safe to compact.
|
|
"""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-1'),
|
|
self._create_request_confirmation_call_event(
|
|
3.0, 'inv2', 'confirm-1', 'call-1'
|
|
),
|
|
self._create_hitl_confirmation_event(4.0, 'inv2', 'call-1'),
|
|
self._create_request_confirmation_response_event(
|
|
5.0, 'inv3', 'confirm-1'
|
|
),
|
|
self._create_function_response_event(6.0, 'inv3', 'call-1'),
|
|
self._create_event(7.0, 'inv4', 'e7'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 7.0, 'Summary resolved real hitl'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
# Both call-1 and confirm-1 are resolved, so the full span compacts.
|
|
self.assertEqual(
|
|
compacted_inv_ids,
|
|
['inv1', 'inv2', 'inv2', 'inv2', 'inv3', 'inv3', 'inv4'],
|
|
)
|
|
|
|
async def test_sliding_window_stops_compaction_at_open_obligations(
|
|
self,
|
|
):
|
|
"""Compaction stops at the first still-open call/HITL obligation."""
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=self.mock_compactor,
|
|
compaction_interval=2,
|
|
overlap_size=0,
|
|
),
|
|
)
|
|
# inv2's call-a only resolves at inv5, and inv4's call-b never resolves,
|
|
# so no prefix past inv1 is self-contained.
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-a'),
|
|
self._create_hitl_confirmation_event(3.0, 'inv3', 'call-a'),
|
|
self._create_function_call_event(4.0, 'inv4', 'call-b'),
|
|
self._create_function_response_event(5.0, 'inv5', 'call-a'),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 3.0, 'Summary including resolved hitl'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
async def test_token_threshold_resolved_hitl_outside_window_is_compactable(
|
|
self,
|
|
):
|
|
"""Token-threshold: HITL with resolver past the truncation point compacts."""
|
|
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,
|
|
),
|
|
)
|
|
events = [
|
|
self._create_event(1.0, 'inv1', 'e1'),
|
|
self._create_function_call_event(2.0, 'inv2', 'call-a'),
|
|
self._create_hitl_confirmation_event(3.0, 'inv3', 'call-a'),
|
|
self._create_function_call_event(4.0, 'inv4', 'call-b'),
|
|
self._create_function_response_event(
|
|
5.0, 'inv5', 'call-a', prompt_token_count=100
|
|
),
|
|
]
|
|
session = Session(app_name='test', user_id='u1', id='s1', events=events)
|
|
|
|
mock_compacted_event = self._create_compacted_event(
|
|
1.0, 3.0, 'Summary including resolved hitl'
|
|
)
|
|
self.mock_compactor.maybe_summarize_events.return_value = (
|
|
mock_compacted_event
|
|
)
|
|
|
|
await self._run_sliding_window(app, session, self.mock_session_service)
|
|
|
|
compacted_events_arg = self.mock_compactor.maybe_summarize_events.call_args[
|
|
1
|
|
]['events']
|
|
compacted_inv_ids = [e.invocation_id for e in compacted_events_arg]
|
|
self.assertEqual(compacted_inv_ids, ['inv1'])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_compaction_for_token_threshold_adds_summary_trace(
|
|
span_exporter: InMemorySpanExporter,
|
|
):
|
|
session = Session(
|
|
app_name='app',
|
|
user_id='user',
|
|
id='session-id',
|
|
events=[
|
|
_create_trace_test_event(
|
|
timestamp=1.0, invocation_id='inv1', text='e1'
|
|
),
|
|
_create_trace_test_event(
|
|
timestamp=2.0, invocation_id='inv2', text='e2'
|
|
),
|
|
_create_trace_test_event(
|
|
timestamp=3.0,
|
|
invocation_id='inv3',
|
|
text='e3',
|
|
prompt_token_count=100,
|
|
),
|
|
],
|
|
)
|
|
session_service = AsyncMock(spec=BaseSessionService)
|
|
compacted_event = _create_trace_compacted_event(
|
|
start_ts=1.0, end_ts=2.0, summary_text='summary'
|
|
)
|
|
summarizer = _StubSummarizer(compacted_event)
|
|
config = EventsCompactionConfig(
|
|
summarizer=summarizer,
|
|
compaction_interval=999,
|
|
overlap_size=0,
|
|
token_threshold=50,
|
|
event_retention_size=1,
|
|
)
|
|
|
|
compacted = (
|
|
await (
|
|
compaction_module._run_compaction_for_token_threshold_config(
|
|
config=config,
|
|
session=session,
|
|
session_service=session_service,
|
|
agent=Mock(spec=BaseAgent),
|
|
)
|
|
)
|
|
)
|
|
|
|
assert compacted is True
|
|
spans = span_exporter.get_finished_spans()
|
|
summary_span = next(
|
|
span for span in spans if span.name == 'compact_events token_threshold'
|
|
)
|
|
assert summary_span.attributes['gen_ai.conversation.id'] == 'session-id'
|
|
assert (
|
|
summary_span.attributes['gen_ai.compaction.trigger'] == 'token_threshold'
|
|
)
|
|
assert summary_span.attributes['gen_ai.compaction.event_count'] == 2
|
|
assert summary_span.attributes['gen_ai.compaction.token_threshold'] == 50
|
|
assert summary_span.attributes['gen_ai.compaction.event_retention_size'] == 1
|
|
assert (
|
|
summary_span.attributes['gen_ai.compaction.result_event_id']
|
|
== 'compacted-event-id'
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_compaction_for_sliding_window_adds_summary_trace(
|
|
span_exporter: InMemorySpanExporter,
|
|
):
|
|
compacted_event = _create_trace_compacted_event(
|
|
start_ts=1.0, end_ts=4.0, summary_text='summary'
|
|
)
|
|
summarizer = _StubSummarizer(compacted_event)
|
|
app = App(
|
|
name='test',
|
|
root_agent=Mock(spec=BaseAgent),
|
|
events_compaction_config=EventsCompactionConfig(
|
|
summarizer=summarizer,
|
|
compaction_interval=2,
|
|
overlap_size=1,
|
|
),
|
|
)
|
|
session = Session(
|
|
app_name='test',
|
|
user_id='u1',
|
|
id='session-id',
|
|
events=[
|
|
_create_trace_test_event(
|
|
timestamp=1.0, invocation_id='inv1', text='e1'
|
|
),
|
|
_create_trace_test_event(
|
|
timestamp=2.0, invocation_id='inv2', text='e2'
|
|
),
|
|
_create_trace_test_event(
|
|
timestamp=3.0, invocation_id='inv3', text='e3'
|
|
),
|
|
_create_trace_test_event(
|
|
timestamp=4.0, invocation_id='inv4', text='e4'
|
|
),
|
|
],
|
|
)
|
|
session_service = AsyncMock(spec=BaseSessionService)
|
|
|
|
async for _ in _run_compaction_for_sliding_window(
|
|
app, session, session_service
|
|
):
|
|
pass
|
|
|
|
spans = span_exporter.get_finished_spans()
|
|
summary_span = next(
|
|
span for span in spans if span.name == 'compact_events sliding_window'
|
|
)
|
|
assert summary_span.attributes['gen_ai.conversation.id'] == 'session-id'
|
|
assert (
|
|
summary_span.attributes['gen_ai.compaction.trigger'] == 'sliding_window'
|
|
)
|
|
assert summary_span.attributes['gen_ai.compaction.event_count'] == 4
|
|
assert summary_span.attributes['gen_ai.compaction.compaction_interval'] == 2
|
|
assert summary_span.attributes['gen_ai.compaction.overlap_size'] == 1
|
|
assert (
|
|
summary_span.attributes['gen_ai.compaction.result_event_id']
|
|
== 'compacted-event-id'
|
|
)
|