feat: Add ability to use existing sessions for streaming_agent_run_with_events calls.

PiperOrigin-RevId: 816470728

Author: vertex-sdk-bot

vertexai/agent_engines/templates/adk.py

@@ -169,6 +169,9 @@ def __init__(self, **kwargs):
         self.user_id: Optional[str] = kwargs.get("user_id", _DEFAULT_USER_ID)
         # The user ID.
 
+        self.session_id: Optional[str] = kwargs.get("session_id")
+        # The session ID.
+
 
 class _StreamingRunResponse:
     """Response object for `streaming_agent_run_with_events` method.
@@ -181,6 +184,8 @@ def __init__(self, **kwargs):
         # List of generated events.
         self.artifacts: Optional[List[_Artifact]] = kwargs.get("artifacts")
         # List of artifacts belonging to the session.
+        self.session_id: Optional[str] = kwargs.get("session_id")
+        # The session ID.
 
     def dump(self) -> Dict[str, Any]:
         from vertexai.agent_engines import _utils
@@ -194,6 +199,8 @@ def dump(self) -> Dict[str, Any]:
                 result["events"].append(event_dict)
         if self.artifacts:
             result["artifacts"] = [artifact.dump() for artifact in self.artifacts]
+        if self.session_id:
+            result["session_id"] = self.session_id
         return result
 
 
@@ -402,7 +409,10 @@ async def _init_session(
                 auth = _Authorization(**auth)
                 session_state[f"temp:{auth_id}"] = auth.access_token
 
-        session_id = f"temp_session_{random.randbytes(8).hex()}"
+        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,
@@ -450,7 +460,9 @@ async def _convert_response_events(
         """Converts the events to the streaming run response object."""
         import collections
 
-        result = _StreamingRunResponse(events=events, artifacts=[])
+        result = _StreamingRunResponse(
+            events=events, artifacts=[], session_id=session_id
+        )
 
         # Save the generated artifacts into the result object.
         artifact_versions = collections.defaultdict(list)
@@ -685,22 +697,35 @@ async def streaming_agent_run_with_events(self, request_json: str):
         request = _StreamRunRequest(**json.loads(request_json))
         if not self._tmpl_attrs.get("in_memory_runner"):
             self.set_up()
-        if not self._tmpl_attrs.get("artifact_service"):
-            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 = await self._init_session(
-            session_service=self._tmpl_attrs.get("in_memory_session_service"),
-            artifact_service=self._tmpl_attrs.get("in_memory_artifact_service"),
-            request=request,
-        )
+        session_service = self._tmpl_attrs.get("in_memory_session_service")
+        artifact_service = self._tmpl_attrs.get("in_memory_artifact_service")
+        # 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=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,
+            )
         if not session:
             raise RuntimeError("Session initialization failed.")
 
-        # Run the agent.
+        # Run the agent
         message_for_agent = types.Content(**request.message)
         try:
             async for event in self._tmpl_attrs.get("in_memory_runner").run_async(
@@ -712,15 +737,16 @@ async def streaming_agent_run_with_events(self, request_json: str):
                     user_id=request.user_id,
                     session_id=session.id,
                     events=[event],
-                    artifact_service=self._tmpl_attrs.get("in_memory_artifact_service"),
+                    artifact_service=artifact_service,
                 )
                 yield converted_event
         finally:
-            await self._tmpl_attrs.get("in_memory_session_service").delete_session(
-                app_name=self._tmpl_attrs.get("app_name"),
-                user_id=request.user_id,
-                session_id=session.id,
-            )
+            if session and not request.session_id:
+                await session_service.delete_session(
+                    app_name=self._tmpl_attrs.get("app_name"),
+                    user_id=request.user_id,
+                    session_id=session.id,
+                )
 
     async def async_get_session(
         self,

vertexai/preview/reasoning_engines/templates/adk.py

@@ -183,6 +183,9 @@ def __init__(self, **kwargs):
         self.user_id: Optional[str] = kwargs.get("user_id", _DEFAULT_USER_ID)
         # The user ID.
 
+        self.session_id: Optional[str] = kwargs.get("session_id")
+        # The session ID.
+
 
 class _StreamingRunResponse:
     """Response object for `streaming_agent_run_with_events` method.
@@ -195,6 +198,8 @@ def __init__(self, **kwargs):
         # List of generated events.
         self.artifacts: Optional[List[_Artifact]] = kwargs.get("artifacts")
         # List of artifacts belonging to the session.
+        self.session_id: Optional[str] = kwargs.get("session_id")
+        # The session ID.
 
     def dump(self) -> Dict[str, Any]:
         from vertexai.agent_engines import _utils
@@ -208,6 +213,8 @@ def dump(self) -> Dict[str, Any]:
                 result["events"].append(event_dict)
         if self.artifacts:
             result["artifacts"] = [artifact.dump() for artifact in self.artifacts]
+        if self.session_id:
+            result["session_id"] = self.session_id
         return result
 
 
@@ -383,7 +390,10 @@ async def _init_session(
                 auth = _Authorization(**auth)
                 session_state[f"temp:{auth_id}"] = auth.access_token
 
-        session_id = f"temp_session_{random.randbytes(8).hex()}"
+        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,
@@ -431,7 +441,9 @@ async def _convert_response_events(
         """Converts the events to the streaming run response object."""
         import collections
 
-        result = _StreamingRunResponse(events=events, artifacts=[])
+        result = _StreamingRunResponse(
+            events=events, artifacts=[], session_id=session_id
+        )
 
         # Save the generated artifacts into the result object.
         artifact_versions = collections.defaultdict(list)
@@ -735,21 +747,33 @@ 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("artifact_service"):
-                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 = await self._init_session(
-                session_service=self._tmpl_attrs.get("in_memory_session_service"),
-                artifact_service=self._tmpl_attrs.get("in_memory_artifact_service"),
-                request=request,
-            )
+            session_service = self._tmpl_attrs.get("in_memory_session_service")
+            artifact_service = self._tmpl_attrs.get("in_memory_artifact_service")
+            # 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=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,
+                )
             if not session:
                 raise RuntimeError("Session initialization failed.")
-
             # Run the agent.
             message_for_agent = types.Content(**request.message)
             try:
@@ -762,17 +786,16 @@ async def _invoke_agent_async():
                         user_id=request.user_id,
                         session_id=session.id,
                         events=[event],
-                        artifact_service=self._tmpl_attrs.get(
-                            "in_memory_artifact_service"
-                        ),
+                        artifact_service=artifact_service,
                     )
                     event_queue.put(converted_event)
             finally:
-                await self._tmpl_attrs.get("in_memory_session_service").delete_session(
-                    app_name=self._tmpl_attrs.get("app_name"),
-                    user_id=request.user_id,
-                    session_id=session.id,
-                )
+                if session and not request.session_id:
+                    await session_service.delete_session(
+                        app_name=self._tmpl_attrs.get("app_name"),
+                        user_id=request.user_id,
+                        session_id=session.id,
+                    )
 
         def _asyncio_thread_main():
             try: