ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
266 lines
9.3 KiB
Python
266 lines
9.3 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 google.adk.apps.llm_event_summarizer import LlmEventSummarizer
|
|
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.models.base_llm import BaseLlm
|
|
from google.adk.models.llm_request import LlmRequest
|
|
from google.adk.models.llm_response import LlmResponse
|
|
from google.genai import types
|
|
from google.genai.types import Content
|
|
from google.genai.types import FunctionCall
|
|
from google.genai.types import FunctionResponse
|
|
from google.genai.types import Part
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'env_variables', ['GOOGLE_AI', 'VERTEX'], indirect=True
|
|
)
|
|
class TestLlmEventSummarizer(unittest.IsolatedAsyncioTestCase):
|
|
|
|
def setUp(self):
|
|
self.mock_llm = AsyncMock(spec=BaseLlm)
|
|
self.mock_llm.model = 'test-model'
|
|
self.compactor = LlmEventSummarizer(llm=self.mock_llm)
|
|
|
|
def _create_event(
|
|
self, timestamp: float, text: str, author: str = 'user'
|
|
) -> Event:
|
|
return Event(
|
|
timestamp=timestamp,
|
|
author=author,
|
|
content=Content(parts=[Part(text=text)]),
|
|
)
|
|
|
|
async def test_maybe_compact_events_success(self):
|
|
events = [
|
|
self._create_event(1.0, 'Hello', 'user'),
|
|
self._create_event(2.0, 'Hi there!', 'model'),
|
|
]
|
|
expected_conversation_history = 'user: Hello\nmodel: Hi there!'
|
|
expected_prompt = self.compactor._DEFAULT_PROMPT_TEMPLATE.format(
|
|
conversation_history=expected_conversation_history
|
|
)
|
|
llm_response = LlmResponse(
|
|
content=Content(parts=[Part(text='Summary')]),
|
|
usage_metadata=None,
|
|
)
|
|
|
|
async def async_gen():
|
|
yield llm_response
|
|
|
|
self.mock_llm.generate_content_async.return_value = async_gen()
|
|
|
|
compacted_event = await self.compactor.maybe_summarize_events(events=events)
|
|
|
|
self.assertIsNotNone(compacted_event)
|
|
self.assertEqual(
|
|
compacted_event.actions.compaction.compacted_content.parts[0].text,
|
|
'Summary',
|
|
)
|
|
self.assertEqual(compacted_event.author, 'user')
|
|
self.assertIsNone(compacted_event.usage_metadata)
|
|
self.assertIsNotNone(compacted_event.actions)
|
|
self.assertIsNotNone(compacted_event.actions.compaction)
|
|
self.assertEqual(compacted_event.actions.compaction.start_timestamp, 1.0)
|
|
self.assertEqual(compacted_event.actions.compaction.end_timestamp, 2.0)
|
|
self.assertEqual(
|
|
compacted_event.actions.compaction.compacted_content.parts[0].text,
|
|
'Summary',
|
|
)
|
|
|
|
self.mock_llm.generate_content_async.assert_called_once()
|
|
args, kwargs = self.mock_llm.generate_content_async.call_args
|
|
llm_request = args[0]
|
|
self.assertIsInstance(llm_request, LlmRequest)
|
|
self.assertEqual(llm_request.model, 'test-model')
|
|
self.assertEqual(llm_request.contents[0].role, 'user')
|
|
self.assertEqual(llm_request.contents[0].parts[0].text, expected_prompt)
|
|
self.assertFalse(kwargs['stream'])
|
|
|
|
async def test_maybe_compact_events_empty_llm_response(self):
|
|
events = [
|
|
self._create_event(1.0, 'Hello', 'user'),
|
|
]
|
|
llm_response = LlmResponse(content=None, usage_metadata=None)
|
|
|
|
async def async_gen():
|
|
yield llm_response
|
|
|
|
self.mock_llm.generate_content_async.return_value = async_gen()
|
|
|
|
compacted_event = await self.compactor.maybe_summarize_events(events=events)
|
|
self.assertIsNone(compacted_event)
|
|
|
|
async def test_maybe_compact_events_includes_usage_metadata(self):
|
|
events = [
|
|
self._create_event(1.0, 'Hello', 'user'),
|
|
self._create_event(2.0, 'Hi there!', 'model'),
|
|
]
|
|
usage_metadata = types.GenerateContentResponseUsageMetadata(
|
|
prompt_token_count=10,
|
|
candidates_token_count=5,
|
|
)
|
|
llm_response = LlmResponse(
|
|
content=Content(parts=[Part(text='Summary')]),
|
|
usage_metadata=usage_metadata,
|
|
)
|
|
|
|
async def async_gen():
|
|
yield llm_response
|
|
|
|
self.mock_llm.generate_content_async.return_value = async_gen()
|
|
|
|
compacted_event = await self.compactor.maybe_summarize_events(events=events)
|
|
|
|
self.assertIsNotNone(compacted_event)
|
|
self.assertEqual(compacted_event.usage_metadata, usage_metadata)
|
|
self.assertEqual(compacted_event.usage_metadata.prompt_token_count, 10)
|
|
self.assertEqual(compacted_event.usage_metadata.candidates_token_count, 5)
|
|
|
|
async def test_maybe_compact_events_empty_input(self):
|
|
compacted_event = await self.compactor.maybe_summarize_events(events=[])
|
|
self.assertIsNone(compacted_event)
|
|
self.mock_llm.generate_content_async.assert_not_called()
|
|
|
|
def test_format_events_for_prompt(self):
|
|
events = [
|
|
self._create_event(1.0, 'User says...', 'user'),
|
|
self._create_event(2.0, 'Model replies...', 'model'),
|
|
self._create_event(3.0, 'Another user input', 'user'),
|
|
self._create_event(4.0, 'More model text', 'model'),
|
|
# Event with no content
|
|
Event(timestamp=5.0, author='user'),
|
|
# Event with empty content part
|
|
Event(
|
|
timestamp=6.0,
|
|
author='model',
|
|
content=Content(parts=[Part(text='')]),
|
|
),
|
|
# Event with function call
|
|
Event(
|
|
timestamp=7.0,
|
|
author='model',
|
|
content=Content(
|
|
parts=[
|
|
Part(
|
|
function_call=FunctionCall(
|
|
id='call_1', name='tool', args={'q': 'x'}
|
|
)
|
|
)
|
|
]
|
|
),
|
|
),
|
|
# Event with function response
|
|
Event(
|
|
timestamp=8.0,
|
|
author='model',
|
|
content=Content(
|
|
parts=[
|
|
Part(
|
|
function_response=FunctionResponse(
|
|
id='call_1',
|
|
name='tool',
|
|
response={'result': 'done'},
|
|
)
|
|
)
|
|
]
|
|
),
|
|
),
|
|
]
|
|
expected_formatted_history = (
|
|
'user: User says...\nmodel: Model replies...\nuser: Another user'
|
|
' input\nmodel: More model text\nmodel called tool:'
|
|
" tool({'q': 'x'})\nTool response from tool: {'result': 'done'}"
|
|
)
|
|
formatted_history = self.compactor._format_events_for_prompt(events)
|
|
self.assertEqual(formatted_history, expected_formatted_history)
|
|
|
|
def test_format_events_for_prompt_includes_thoughts(self):
|
|
events = [
|
|
self._create_event(1.0, 'What is the weather?', 'user'),
|
|
Event(
|
|
timestamp=2.0,
|
|
author='model',
|
|
content=Content(
|
|
parts=[
|
|
Part(text='Let me check the tool output.', thought=True),
|
|
Part(text='It is sunny.'),
|
|
]
|
|
),
|
|
),
|
|
]
|
|
expected_formatted_history = (
|
|
'user: What is the weather?\nmodel (thought): Let me check the tool'
|
|
' output.\nmodel: It is sunny.'
|
|
)
|
|
formatted_history = self.compactor._format_events_for_prompt(events)
|
|
self.assertEqual(formatted_history, expected_formatted_history)
|
|
|
|
def test_format_events_for_prompt_skips_compaction_event_thought(self):
|
|
events = [
|
|
Event(
|
|
timestamp=1.0,
|
|
author='model',
|
|
content=Content(
|
|
parts=[
|
|
Part(text='Stale summarizer reasoning.', thought=True),
|
|
Part(text='Prior summary.'),
|
|
]
|
|
),
|
|
actions=EventActions(
|
|
compaction=EventCompaction(
|
|
start_timestamp=0.0,
|
|
end_timestamp=1.0,
|
|
compacted_content=Content(parts=[Part(text='Prior')]),
|
|
)
|
|
),
|
|
),
|
|
self._create_event(2.0, 'New user input', 'user'),
|
|
]
|
|
expected_formatted_history = 'model: Prior summary.\nuser: New user input'
|
|
formatted_history = self.compactor._format_events_for_prompt(events)
|
|
self.assertEqual(formatted_history, expected_formatted_history)
|
|
|
|
def test_format_events_for_prompt_truncates_large_tool_response(self):
|
|
limit = self.compactor._MAX_TOOL_CONTENT_CHARS
|
|
large_value = 'x' * (limit + 500)
|
|
events = [
|
|
Event(
|
|
timestamp=1.0,
|
|
author='model',
|
|
content=Content(
|
|
parts=[
|
|
Part(
|
|
function_response=FunctionResponse(
|
|
id='call_1',
|
|
name='search',
|
|
response={'data': large_value},
|
|
)
|
|
)
|
|
]
|
|
),
|
|
),
|
|
]
|
|
formatted_history = self.compactor._format_events_for_prompt(events)
|
|
self.assertIn('Tool response from search:', formatted_history)
|
|
self.assertIn('... [truncated', formatted_history)
|
|
self.assertLess(len(formatted_history), len(large_value))
|