feat: Add support for app input in AdkApp template

PiperOrigin-RevId: 825039994

Author: yeesian

vertexai/agent_engines/templates/adk.py

@@ -37,6 +37,13 @@
     except (ImportError, AttributeError):
         Event = Any
 
+    try:
+        from google.adk.apps import App
+
+        App = App
+    except (ImportError, AttributeError):
+        App = Any
+
     try:
         from google.adk.agents import BaseAgent
 
@@ -449,7 +456,8 @@ class AdkApp:
     def __init__(
         self,
         *,
-        agent: "BaseAgent",
+        agent: "BaseAgent" = None,
+        app: "App" = None,
         app_name: Optional[str] = None,
         plugins: Optional[List["BasePlugin"]] = None,
         enable_tracing: Optional[bool] = None,
@@ -505,10 +513,28 @@ def __init__(
             )
             raise ValueError(msg)
 
+        if not agent and not app:
+            raise ValueError("One of `agent` or `app` must be provided.")
+        if app:
+            if app_name:
+                raise ValueError(
+                    "When app is provided, app_name should not be provided."
+                )
+            if agent:
+                raise ValueError(
+                    "When app is provided, agent should not be provided."
+                )
+            if plugins:
+                raise ValueError(
+                    "When app is provided, plugins should not be provided and"
+                    " should be provided in the app instead."
+                )
+
         self._tmpl_attrs: Dict[str, Any] = {
             "project": initializer.global_config.project,
             "location": initializer.global_config.location,
             "agent": agent,
+            "app": app,
             "app_name": app_name,
             "plugins": plugins,
             "enable_tracing": enable_tracing,
@@ -625,6 +651,7 @@ def clone(self):
 
         return self.__class__(
             agent=copy.deepcopy(self._tmpl_attrs.get("agent")),
+            app=copy.deepcopy(self._tmpl_attrs.get("app")),
             enable_tracing=self._tmpl_attrs.get("enable_tracing"),
             app_name=self._tmpl_attrs.get("app_name"),
             plugins=self._tmpl_attrs.get("plugins"),
@@ -775,18 +802,24 @@ def tracing_enabled() -> bool:
 
         self._tmpl_attrs["runner"] = Runner(
             agent=self._tmpl_attrs.get("agent"),
+            app=self._tmpl_attrs.get("app"),
             plugins=self._tmpl_attrs.get("plugins"),
             session_service=self._tmpl_attrs.get("session_service"),
             artifact_service=self._tmpl_attrs.get("artifact_service"),
             memory_service=self._tmpl_attrs.get("memory_service"),
-            app_name=self._tmpl_attrs.get("app_name"),
+            app_name=(
+                None if self._tmpl_attrs.get("app") else self._tmpl_attrs.get("app_name")
+            ),
         )
         self._tmpl_attrs["in_memory_session_service"] = InMemorySessionService()
         self._tmpl_attrs["in_memory_artifact_service"] = InMemoryArtifactService()
         self._tmpl_attrs["in_memory_memory_service"] = InMemoryMemoryService()
         self._tmpl_attrs["in_memory_runner"] = Runner(
-            app_name=self._tmpl_attrs.get("app_name"),
+            app_name=(
+                None if self._tmpl_attrs.get("app") else self._tmpl_attrs.get("app_name")
+            ),
             agent=self._tmpl_attrs.get("agent"),
+            app=self._tmpl_attrs.get("app"),
             plugins=self._tmpl_attrs.get("plugins"),
             session_service=self._tmpl_attrs.get("in_memory_session_service"),
             artifact_service=self._tmpl_attrs.get("in_memory_artifact_service"),
@@ -968,12 +1001,13 @@ async def streaming_agent_run_with_events(self, request_json: str):
             self.set_up()
         session_service = self._tmpl_attrs.get("in_memory_session_service")
         artifact_service = self._tmpl_attrs.get("in_memory_artifact_service")
+        app = self._tmpl_attrs.get("app")
         # 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"),
+                    app_name=app.name if app else self._tmpl_attrs.get("app_name"),
                     user_id=request.user_id,
                     session_id=request.session_id,
                 )
@@ -1006,8 +1040,9 @@ async def streaming_agent_run_with_events(self, request_json: str):
                 yield converted_event
         finally:
             if session and not request.session_id:
+                app = self._tmpl_attrs.get("app")
                 await session_service.delete_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,
                     session_id=session.id,
                 )
@@ -1039,8 +1074,9 @@ async def async_get_session(
         """
         if not self._tmpl_attrs.get("session_service"):
             self.set_up()
+        app = self._tmpl_attrs.get("app")
         session = await self._tmpl_attrs.get("session_service").get_session(
-            app_name=self._tmpl_attrs.get("app_name"),
+            app_name=app.name if app else self._tmpl_attrs.get("app_name"),
             user_id=user_id,
             session_id=session_id,
             **kwargs,
@@ -1116,8 +1152,9 @@ async def async_list_sessions(self, *, user_id: str, **kwargs):
         """
         if not self._tmpl_attrs.get("session_service"):
             self.set_up()
+        app = self._tmpl_attrs.get("app")
         return await self._tmpl_attrs.get("session_service").list_sessions(
-            app_name=self._tmpl_attrs.get("app_name"),
+            app_name=app.name if app else self._tmpl_attrs.get("app_name"),
             user_id=user_id,
             **kwargs,
         )
@@ -1188,8 +1225,9 @@ async def async_create_session(
         """
         if not self._tmpl_attrs.get("session_service"):
             self.set_up()
+        app = self._tmpl_attrs.get("app")
         session = await self._tmpl_attrs.get("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=user_id,
             session_id=session_id,
             state=state,
@@ -1269,8 +1307,9 @@ async def async_delete_session(
         """
         if not self._tmpl_attrs.get("session_service"):
             self.set_up()
+        app = self._tmpl_attrs.get("app")
         await self._tmpl_attrs.get("session_service").delete_session(
-            app_name=self._tmpl_attrs.get("app_name"),
+            app_name=app.name if app else self._tmpl_attrs.get("app_name"),
             user_id=user_id,
             session_id=session_id,
             **kwargs,
@@ -1359,8 +1398,9 @@ async def async_search_memory(self, *, user_id: str, query: str):
         """
         if not self._tmpl_attrs.get("memory_service"):
             self.set_up()
+        app = self._tmpl_attrs.get("app")
         return await self._tmpl_attrs.get("memory_service").search_memory(
-            app_name=self._tmpl_attrs.get("app_name"),
+            app_name=app.name if app else self._tmpl_attrs.get("app_name"),
             user_id=user_id,
             query=query,
         )