From c4acba0904fb17030da10afcb18943de9bc3c2ee Mon Sep 17 00:00:00 2001 From: Tongzhou Jiang Date: Thu, 8 Jan 2026 15:37:14 -0800 Subject: [PATCH] chore: make sure memory service is setup in `streaming_agent_run_with_response` PiperOrigin-RevId: 853909960 --- vertexai/agent_engines/templates/adk.py | 25 +++++++++--------- .../reasoning_engines/templates/adk.py | 26 +++++++++---------- 2 files changed, 26 insertions(+), 25 deletions(-) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index e717410fe6..cf5e005767 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -1149,18 +1149,19 @@ async def streaming_agent_run_with_events(self, request_json: str): from google.genai.errors import ClientError request = _StreamRunRequest(**json.loads(request_json)) - if not self._tmpl_attrs.get("in_memory_runner"): - self.set_up() - if not self._tmpl_attrs.get("runner"): - self.set_up() - # Prepare the in-memory session. - if not self._tmpl_attrs.get("in_memory_artifact_service"): - self.set_up() - if not self._tmpl_attrs.get("artifact_service"): - self.set_up() - if not self._tmpl_attrs.get("in_memory_session_service"): - self.set_up() - if not self._tmpl_attrs.get("session_service"): + if not any( + self._tmpl_attrs.get(service) + for service in ( + "in_memory_runner", + "runner", + "in_memory_artifact_service", + "artifact_service", + "in_memory_session_service", + "session_service", + "in_memory_memory_service", + "memory_service", + ) + ): self.set_up() app = self._tmpl_attrs.get("app") diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 872de508ef..6c67401f2f 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -989,20 +989,20 @@ def streaming_agent_run_with_events(self, request_json: str): async def _invoke_agent_async(): request = _StreamRunRequest(**json.loads(request_json)) - if not self._tmpl_attrs.get("in_memory_runner"): - self.set_up() - if not self._tmpl_attrs.get("runner"): - self.set_up() - # Prepare the in-memory session. - if not self._tmpl_attrs.get("in_memory_artifact_service"): - self.set_up() - if not self._tmpl_attrs.get("artifact_service"): - self.set_up() - if not self._tmpl_attrs.get("in_memory_session_service"): - self.set_up() - if not self._tmpl_attrs.get("session_service"): + if not any( + self._tmpl_attrs.get(service) + for service in ( + "in_memory_runner", + "runner", + "in_memory_artifact_service", + "artifact_service", + "in_memory_session_service", + "session_service", + "in_memory_memory_service", + "memory_service", + ) + ): self.set_up() - # Try to get the session, if it doesn't exist, create a new one. if request.session_id: session_service = self._tmpl_attrs.get("session_service")