fix: investigate save artifact in init_session

PiperOrigin-RevId: 833983514

Author: Tongzhou-Jiang

vertexai/agent_engines/templates/adk.py

@@ -135,8 +135,10 @@ def is_version_sufficient(version_to_check: str) -> bool:
 
 class _ArtifactVersion:
     def __init__(self, **kwargs):
+        from google.genai import types
+
         self.version: Optional[str] = kwargs.get("version")
-        self.data = kwargs.get("data")
+        self.data: Optional[types.Part] = kwargs.get("data")
 
     def dump(self) -> Dict[str, Any]:
         result = {}
@@ -603,15 +605,20 @@ async def _init_session(
         """Initializes the session, and returns the session id."""
         from google.adk.events.event import Event
 
+        from google.cloud.aiplatform import base
+
+        _LOGGER = base.Logger(__name__)
+
         session_state = None
         if request.authorizations:
             session_state = {}
             for auth_id, auth in request.authorizations.items():
                 auth = _Authorization(**auth)
                 session_state[auth_id] = auth.access_token
 
+        app = self._tmpl_attrs.get("app")
         session = await session_service.create_session(
-            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,
             state=session_state,
         )
@@ -627,8 +634,9 @@ async def _init_session(
                     artifact.versions, key=lambda x: x["version"]
                 ):
                     version_data = _ArtifactVersion(**version_data)
+                    _LOGGER.info(f'Saving artifact {version_data.data}, type {type(version_data.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,
                         filename=artifact.file_name,