fix: save artifact in streaming agent run with events when multiturn

PiperOrigin-RevId: 873078156

Author: Tongzhou-Jiang

vertexai/agent_engines/templates/adk.py

@@ -683,6 +683,20 @@ async def _init_session(
         if request.events:
             for event in request.events:
                 await session_service.append_event(session, Event(**event))
+        if request.artifacts:
+            await self._save_artifacts(
+                session.id, artifact_service, request
+            )
+        return session
+
+    async def _save_artifacts(
+        self,
+        session_id: str,
+        artifact_service: "BaseArtifactService",
+        request: _StreamRunRequest,
+    ):
+        """Saves the artifacts."""
+        app = self._tmpl_attrs.get("app")
         if request.artifacts:
             for artifact in request.artifacts:
                 artifact = _Artifact(**artifact)
@@ -693,7 +707,7 @@ async def _init_session(
                     saved_version = await artifact_service.save_artifact(
                         app_name=app.name if app else self._tmpl_attrs.get("app_name"),
                         user_id=request.user_id,
-                        session_id=session.id,
+                        session_id=session_id,
                         filename=artifact.file_name,
                         artifact=version_data.data,
                     )
@@ -707,7 +721,6 @@ async def _init_session(
                             saved_version,
                             version_data.version,
                         )
-        return session
 
     async def _convert_response_events(
         self,
@@ -1209,6 +1222,11 @@ async def streaming_agent_run_with_events(self, request_json: str):
                     user_id=request.user_id,
                     session_id=request.session_id,
                 )
+                await self._save_artifacts(
+                        session_id=request.session_id,
+                        artifact_service=artifact_service,
+                        request=request,
+                )
             except ClientError:
                 pass
             if not session:

vertexai/preview/reasoning_engines/templates/adk.py

@@ -616,6 +616,23 @@ async def _init_session(
         if request.events:
             for event in request.events:
                 await session_service.append_event(session, Event(**event))
+        if request.artifacts:
+            await self._save_artifacts(
+                session_id=session.id,
+                artifact_service=artifact_service,
+                request=request,
+            )
+
+        return session
+
+    async def _save_artifacts(
+        self,
+        session_id: str,
+        artifact_service: "BaseArtifactService",
+        request: _StreamRunRequest,
+    ):
+        """Saves the artifacts."""
+        app = self._tmpl_attrs.get("app")
         if request.artifacts:
             for artifact in request.artifacts:
                 artifact = _Artifact(**artifact)
@@ -624,9 +641,9 @@ async def _init_session(
                 ):
                     version_data = _ArtifactVersion(**version_data)
                     saved_version = await artifact_service.save_artifact(
-                        app_name=self._tmpl_attrs.get("app_name"),
+                        app_name=app.name if app else self._tmpl_attrs.get("app_name"),
                         user_id=request.user_id,
-                        session_id=session.id,
+                        session_id=session_id,
                         filename=artifact.file_name,
                         artifact=version_data.data,
                     )
@@ -640,7 +657,6 @@ async def _init_session(
                             saved_version,
                             version_data.version,
                         )
-        return session
 
     async def _convert_response_events(
         self,
@@ -1054,6 +1070,11 @@ async def _invoke_agent_async():
                         artifact_service=artifact_service,
                         request=request,
                     )
+                    await self._save_artifacts(
+                        session_id=request.session_id,
+                        artifact_service=artifact_service,
+                        request=request,
+                    )
             else:
                 # Not providing a session ID will create a new in-memory session.
                 session_service = self._tmpl_attrs.get("in_memory_session_service")