feat: enable persistent short-term memory to pass response-id

Author: zakahan

veadk/agent.py

@@ -29,7 +29,7 @@
 from google.adk.agents.context_cache_config import ContextCacheConfig
 from google.adk.agents.llm_agent import InstructionProvider, ToolUnion
 from google.adk.agents.run_config import StreamingMode
-from google.adk.events import Event
+from google.adk.events import Event, EventActions
 from google.adk.models.lite_llm import LiteLlm
 from google.adk.runners import Runner
 from google.genai import types
@@ -240,6 +240,18 @@ async def _run_async_impl(
 
         async for event in super()._run_async_impl(ctx):
             yield event
+            if self.enable_responses and event.cache_metadata:
+                # for persistent short-term memory with response api
+                session_state_event = Event(
+                    invocation_id=event.invocation_id,
+                    author=event.author,
+                    actions=EventActions(
+                        state_delta={
+                            "response_id": event.cache_metadata.cache_name,
+                        }
+                    ),
+                )
+                yield session_state_event
 
     async def _run(
         self,

veadk/memory/short_term_memory.py

@@ -32,6 +32,7 @@
 from veadk.memory.short_term_memory_backends.sqlite_backend import (
     SQLiteSTMBackend,
 )
+from veadk.models.ark_transform import build_cache_metadata
 from veadk.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -41,14 +42,29 @@ def wrap_get_session_with_callbacks(obj, callback_fn: Callable):
     get_session_fn = getattr(obj, "get_session")
 
     @wraps(get_session_fn)
-    def wrapper(*args, **kwargs):
-        result = get_session_fn(*args, **kwargs)
+    async def wrapper(*args, **kwargs):
+        result = await get_session_fn(*args, **kwargs)
         callback_fn(result, *args, **kwargs)
         return result
 
     setattr(obj, "get_session", wrapper)
 
 
+def enable_responses_api_for_session_service(result, *args, **kwargs):
+    if result and isinstance(result, Session):
+        if result.events:
+            for event in result.events:
+                if (
+                    event.actions
+                    and event.actions.state_delta
+                    and not event.cache_metadata
+                    and "response_id" in event.actions.state_delta
+                ):
+                    event.cache_metadata = build_cache_metadata(
+                        response_id=event.actions.state_delta.get("response_id"),
+                    )
+
+
 class ShortTermMemory(BaseModel):
     """Short term memory for agent execution.
 
@@ -170,6 +186,11 @@ def model_post_init(self, __context: Any) -> None:
                         db_kwargs=self.db_kwargs, **self.backend_configs
                     ).session_service
 
+        if self.backend != "local":
+            wrap_get_session_with_callbacks(
+                self._session_service, enable_responses_api_for_session_service
+            )
+
         if self.after_load_memory_callback:
             wrap_get_session_with_callbacks(
                 self._session_service, self.after_load_memory_callback