diff --git a/src/google/adk/apps/llm_event_summarizer.py b/src/google/adk/apps/llm_event_summarizer.py index 8081dd4f13f..9916c0cace0 100644 --- a/src/google/adk/apps/llm_event_summarizer.py +++ b/src/google/adk/apps/llm_event_summarizer.py @@ -13,7 +13,10 @@ # limitations under the License. from __future__ import annotations +import logging + from google.genai.types import Content +from google.genai.types import FinishReason from google.genai.types import Part from ..apps.base_events_summarizer import BaseEventsSummarizer @@ -24,6 +27,16 @@ from ..models.llm_request import LlmRequest from ..utils.context_utils import Aclosing +logger = logging.getLogger('google_adk.' + __name__) + +# A summary that ended for any other reason (e.g. MAX_TOKENS or SAFETY) is +# incomplete and must not replace the events it was meant to summarize. +_COMPLETE_FINISH_REASONS = ( + None, + FinishReason.STOP, + FinishReason.FINISH_REASON_UNSPECIFIED, +) + class LlmEventSummarizer(BaseEventsSummarizer): """An LLM-based event summarizer for sliding window compaction. @@ -150,6 +163,7 @@ async def maybe_summarize_events( ) summary_content = None summary_usage_metadata = None + summary_finish_reason = None async with Aclosing( self._llm.generate_content_async(llm_request, stream=False) ) as agen: @@ -157,11 +171,20 @@ async def maybe_summarize_events( if llm_response.content: summary_content = llm_response.content summary_usage_metadata = llm_response.usage_metadata + summary_finish_reason = llm_response.finish_reason break if summary_content is None: return None + if summary_finish_reason not in _COMPLETE_FINISH_REASONS: + logger.warning( + 'Discarding summary that ended with finish reason %s; keeping the' + ' uncompacted history.', + summary_finish_reason, + ) + return None + # Ensure the compacted content has the role 'model' summary_content.role = 'model' diff --git a/tests/unittests/apps/test_llm_event_summarizer.py b/tests/unittests/apps/test_llm_event_summarizer.py index 2442167bb18..6d805154fb3 100644 --- a/tests/unittests/apps/test_llm_event_summarizer.py +++ b/tests/unittests/apps/test_llm_event_summarizer.py @@ -118,6 +118,64 @@ async def async_gen(): compacted_event = await self.compactor.maybe_summarize_events(events=events) self.assertIsNone(compacted_event) + async def test_maybe_compact_events_rejects_incomplete_summary(self): + events = [ + self._create_event(1.0, 'The delivery address is 123 Maple Avenue.'), + self._create_event(2.0, 'Saved.', 'model'), + ] + for finish_reason in ( + types.FinishReason.MAX_TOKENS, + types.FinishReason.SAFETY, + types.FinishReason.OTHER, + ): + with self.subTest(finish_reason=finish_reason): + llm_response = LlmResponse( + content=Content(parts=[Part(text='The delivery address is')]), + finish_reason=finish_reason, + ) + + async def async_gen(response=llm_response): + yield 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_accepts_complete_summary(self): + events = [ + self._create_event(1.0, 'Hello', 'user'), + self._create_event(2.0, 'Hi there!', 'model'), + ] + for finish_reason in ( + None, + types.FinishReason.STOP, + types.FinishReason.FINISH_REASON_UNSPECIFIED, + ): + with self.subTest(finish_reason=finish_reason): + llm_response = LlmResponse( + content=Content(parts=[Part(text='Summary')]), + finish_reason=finish_reason, + ) + + async def async_gen(response=llm_response): + yield 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', + ) + async def test_maybe_compact_events_includes_usage_metadata(self): events = [ self._create_event(1.0, 'Hello', 'user'),