From 7e4c80ce344e761f9c6ef1b779b74562eca21eb3 Mon Sep 17 00:00:00 2001 From: Giselle van Dongen Date: Thu, 12 Mar 2026 19:08:48 +0700 Subject: [PATCH 1/3] Flush compaction events in Google ADK integration --- python/restate/ext/adk/session.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/restate/ext/adk/session.py b/python/restate/ext/adk/session.py index 2b25346..3b52f1e 100644 --- a/python/restate/ext/adk/session.py +++ b/python/restate/ext/adk/session.py @@ -93,10 +93,13 @@ async def append_event(self, session: Session, event: Event) -> Event: """Appends an event to a session object.""" if event.partial: return event - # For now, we also store temp state event = self._trim_temp_delta_state(event) self._update_session_state(session, event) session.events.append(event) + # Compaction runs after after_run_callback (which flushes the session), + # so compaction events must be flushed explicitly here. + if event.actions and event.actions.compaction: + await self.flush_session_state(session) return event async def flush_session_state(self, session: Session): From faf61c0a6d2df82b810e5702fcb5338b624da2f4 Mon Sep 17 00:00:00 2001 From: Giselle van Dongen Date: Thu, 12 Mar 2026 20:03:52 +0700 Subject: [PATCH 2/3] Add ctx.run for LLM summarizer step --- python/restate/ext/adk/__init__.py | 2 + python/restate/ext/adk/summarizer.py | 72 ++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 python/restate/ext/adk/summarizer.py diff --git a/python/restate/ext/adk/__init__.py b/python/restate/ext/adk/__init__.py index f86a973..eb03175 100644 --- a/python/restate/ext/adk/__init__.py +++ b/python/restate/ext/adk/__init__.py @@ -12,6 +12,7 @@ from .session import RestateSessionService from .plugin import RestatePlugin +from .summarizer import RestateEventSummarizer from restate import ObjectContext, Context from restate.extensions import current_context @@ -35,6 +36,7 @@ def restate_context() -> Context: __all__ = [ "RestateSessionService", "RestatePlugin", + "RestateEventSummarizer", "restate_object_context", "restate_context", ] diff --git a/python/restate/ext/adk/summarizer.py b/python/restate/ext/adk/summarizer.py new file mode 100644 index 0000000..a70507c --- /dev/null +++ b/python/restate/ext/adk/summarizer.py @@ -0,0 +1,72 @@ +# +# Copyright (c) 2023-2025 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Restate-aware event summarizer for ADK compaction. + +Wraps the LlmEventSummarizer so the summarization call is journaled +through ctx.run, making it deterministic on replay. +""" + +import restate + +from datetime import timedelta +from typing import Optional + +from google.adk.apps.base_events_summarizer import BaseEventsSummarizer +from google.adk.apps.llm_event_summarizer import LlmEventSummarizer +from google.adk.events.event import Event +from google.adk.models.base_llm import BaseLlm + +from restate.extensions import current_context + + +class RestateEventSummarizer(BaseEventsSummarizer): + """Event summarizer that journals the LLM call through Restate ctx.run. + + Creates an LlmEventSummarizer internally and wraps its call in ctx.run_typed + so the result is persisted in the Restate journal and replayed deterministically. + """ + + def __init__( + self, + llm: BaseLlm, + prompt_template: Optional[str] = None, + max_retries: int = 10, + ): + self._inner = LlmEventSummarizer(llm=llm, prompt_template=prompt_template) + self._max_retries = max_retries + + async def maybe_summarize_events( + self, *, events: list[Event] + ) -> Optional[Event]: + if not events: + return None + + ctx = current_context() + if ctx is None: + raise RuntimeError( + "No Restate context found. " + "RestateEventSummarizer must be used from within a Restate handler." + ) + + inner = self._inner + + async def call_inner() -> Optional[Event]: + return await inner.maybe_summarize_events(events=events) + + return await ctx.run_typed( + "compaction LLM call", + call_inner, + restate.RunOptions( + max_attempts=self._max_retries, + initial_retry_interval=timedelta(seconds=1), + ), + ) \ No newline at end of file From 70cc997fdc9732f97e27a532e2d23aaabfcb7972 Mon Sep 17 00:00:00 2001 From: Giselle van Dongen Date: Fri, 13 Mar 2026 22:52:29 +0700 Subject: [PATCH 3/3] add factory methods to create from LLM or from custom summarizer --- python/restate/ext/adk/summarizer.py | 35 +++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/python/restate/ext/adk/summarizer.py b/python/restate/ext/adk/summarizer.py index a70507c..c50f17a 100644 --- a/python/restate/ext/adk/summarizer.py +++ b/python/restate/ext/adk/summarizer.py @@ -29,21 +29,44 @@ class RestateEventSummarizer(BaseEventsSummarizer): - """Event summarizer that journals the LLM call through Restate ctx.run. + """Event summarizer that journals the summarization call through Restate ctx.run. - Creates an LlmEventSummarizer internally and wraps its call in ctx.run_typed - so the result is persisted in the Restate journal and replayed deterministically. + Wraps any BaseEventsSummarizer in ctx.run_typed so the result is persisted + in the Restate journal and replayed deterministically. + + Use the factory methods to create instances: + - ``RestateEventSummarizer.from_llm(llm)`` for the default LlmEventSummarizer + - ``RestateEventSummarizer.from_summarizer(summarizer)`` for a custom summarizer """ def __init__( self, - llm: BaseLlm, - prompt_template: Optional[str] = None, + inner: BaseEventsSummarizer, max_retries: int = 10, ): - self._inner = LlmEventSummarizer(llm=llm, prompt_template=prompt_template) + self._inner = inner self._max_retries = max_retries + @staticmethod + def from_llm( + llm: BaseLlm, + prompt_template: Optional[str] = None, + max_retries: int = 10, + ) -> "RestateEventSummarizer": + """Create a RestateEventSummarizer using the default LlmEventSummarizer.""" + return RestateEventSummarizer( + LlmEventSummarizer(llm=llm, prompt_template=prompt_template), + max_retries=max_retries, + ) + + @staticmethod + def from_summarizer( + summarizer: BaseEventsSummarizer, + max_retries: int = 10, + ) -> "RestateEventSummarizer": + """Create a RestateEventSummarizer wrapping a custom summarizer.""" + return RestateEventSummarizer(summarizer, max_retries=max_retries) + async def maybe_summarize_events( self, *, events: list[Event] ) -> Optional[Event]: