feat: Alow VertexAiSession for streaming_agent_run_with_events

PiperOrigin-RevId: 824600367

Author: vertex-sdk-bot

vertexai/agent_engines/templates/adk.py

@@ -559,14 +559,9 @@ async def _init_session(
                 auth = _Authorization(**auth)
                 session_state[f"temp:{auth_id}"] = auth.access_token
 
-        if request.session_id:
-            session_id = request.session_id
-        else:
-            session_id = f"temp_session_{random.randbytes(8).hex()}"
         session = await session_service.create_session(
             app_name=self._tmpl_attrs.get("app_name"),
             user_id=request.user_id,
-            session_id=session_id,
             state=session_state,
         )
         if not session:
@@ -1012,43 +1007,64 @@ async def streaming_agent_run_with_events(self, request_json: str):
 
         import json
         from google.genai import types
+        from google.cloud.aiplatform import base
+        from google.api_core.exceptions import FailedPrecondition
+
+        _LOGGER = base.Logger(__name__)
 
         request = _StreamRunRequest(**json.loads(request_json))
+        _LOGGER.info("request: %s, type: %s", request, type(request))
         if not self._tmpl_attrs.get("in_memory_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("in_memory_session_service"):
             self.set_up()
-        session_service = self._tmpl_attrs.get("in_memory_session_service")
-        artifact_service = self._tmpl_attrs.get("in_memory_artifact_service")
+        session_service = self._tmpl_attrs.get("session_service")
+        artifact_service = self._tmpl_attrs.get("artifact_service")
         app = self._tmpl_attrs.get("app")
+        runner = self._tmpl_attrs.get("runner")
         # Try to get the session, if it doesn't exist, create a new one.
-        session = None
         if request.session_id:
             try:
                 session = await session_service.get_session(
                     app_name=app.name if app else self._tmpl_attrs.get("app_name"),
                     user_id=request.user_id,
                     session_id=request.session_id,
                 )
-            except RuntimeError:
-                pass
-        if not session:
-            #  Fall back to create session if the session is not found.
-            session = await self._init_session(
-                session_service=session_service,
-                artifact_service=artifact_service,
-                request=request,
+            except FailedPrecondition:
+                _LOGGER.info("Session not found, creating a new one.")
+                #  Fall back to create session if the session is not found.
+                session = await self._init_session(
+                    session_service=session_service,
+                    artifact_service=artifact_service,
+                    request=request,
+                )
+            except Exception as e:
+                _LOGGER.error("Failed to get session: %s, type: %s", e, type(e))
+                raise e
+
+        else:
+            # Not providing a session ID will create a new in-memory session.
+            session_service = self._tmpl_attrs.get("in_memory_session_service")
+            artifact_service = self._tmpl_attrs.get("in_memory_artifact_service")
+            runner = self._tmpl_attrs.get("in_memory_runner")
+            session = await session_service.create_session(
+                app_name=self._tmpl_attrs.get("app_name"),
+                user_id=request.user_id,
+                session_id=request.session_id,
             )
         if not session:
             raise RuntimeError("Session initialization failed.")
 
+        _LOGGER.info("session: %s, type: %s", session, type(session))
+        _LOGGER.info("session.id: %s, type: %s", session.id, type(session.id))
+
         # Run the agent
         message_for_agent = types.Content(**request.message)
         try:
-            async for event in self._tmpl_attrs.get("in_memory_runner").run_async(
+            async for event in runner.run_async(
                 user_id=request.user_id,
                 session_id=session.id,
                 new_message=message_for_agent,