feat: Add support for Vertex Express Mode API key in AdkApp

PiperOrigin-RevId: 825638989

Author: vertex-sdk-bot

tests/unit/vertex_adk/test_agent_engine_templates_adk.py

@@ -54,6 +54,7 @@ def __init__(self, name: str, model: str):
 
 _TEST_LOCATION = "us-central1"
 _TEST_PROJECT = "test-project"
+_TEST_API_KEY = "test-api-key"
 _TEST_MODEL = "gemini-2.0-flash"
 _TEST_USER_ID = "test_user_id"
 _TEST_AGENT_NAME = "test_agent"
@@ -761,6 +762,25 @@ def test_dump_event_for_json():
     assert base64.b64decode(part["thought_signature"]) == raw_signature
 
 
+def test_adk_app_initialization_with_api_key():
+    importlib.reload(initializer)
+    importlib.reload(vertexai)
+    try:
+        vertexai.init(api_key=_TEST_API_KEY)
+        app = agent_engines.AdkApp(agent=_TEST_AGENT)
+        assert app._tmpl_attrs.get("project") is None
+        assert app._tmpl_attrs.get("location") is None
+        assert app._tmpl_attrs.get("express_mode_api_key") == _TEST_API_KEY
+        assert app._tmpl_attrs.get("runner") is None
+        app.set_up()
+        assert app._tmpl_attrs.get("runner") is not None
+        assert os.environ.get("GOOGLE_API_KEY") == _TEST_API_KEY
+        assert "GOOGLE_CLOUD_LOCATION" not in os.environ
+        assert "GOOGLE_CLOUD_PROJECT" not in os.environ
+    finally:
+        initializer.global_pool.shutdown(wait=True)
+
+
 @pytest.mark.usefixtures("mock_adk_version")
 class TestAdkAppErrors:
     @pytest.mark.asyncio

vertexai/agent_engines/templates/adk.py

@@ -521,6 +521,7 @@ def __init__(
                 If not provided, a default instrumentor builder will be used.
                 This parameter is ignored if `enable_tracing` is False.
         """
+        import os
         from google.cloud.aiplatform import initializer
 
         adk_version = get_adk_version()
@@ -558,6 +559,7 @@ def __init__(
             "artifact_service_builder": artifact_service_builder,
             "memory_service_builder": memory_service_builder,
             "instrumentor_builder": instrumentor_builder,
+            "express_mode_api_key": initializer.global_config.api_key,
         }
 
     async def _init_session(
@@ -701,9 +703,18 @@ def set_up(self):
 
         os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1"
         project = self._tmpl_attrs.get("project")
-        os.environ["GOOGLE_CLOUD_PROJECT"] = project
+        if project:
+            os.environ["GOOGLE_CLOUD_PROJECT"] = project
         location = self._tmpl_attrs.get("location")
-        os.environ["GOOGLE_CLOUD_LOCATION"] = location
+        if location:
+            os.environ["GOOGLE_CLOUD_LOCATION"] = location
+        express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key")
+        if express_mode_api_key and not project:
+            os.environ["GOOGLE_API_KEY"] = express_mode_api_key
+            # Clear location and project env vars if express mode api key is provided.
+            os.environ.pop("GOOGLE_CLOUD_LOCATION")
+            os.environ.pop("GOOGLE_CLOUD_PROJECT")
+            location = None
 
         # Disable content capture in custom ADK spans unless user enabled
         # tracing explicitly with the old flag
@@ -750,21 +761,57 @@ def set_up(self):
                     VertexAiSessionService,
                 )
 
-                self._tmpl_attrs["session_service"] = VertexAiSessionService(
-                    project=project,
-                    location=location,
-                    agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
-                )
+                if is_version_sufficient("1.18.0"):
+                    if express_mode_api_key:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            express_mode_api_key=express_mode_api_key,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                    else:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            project=project,
+                            location=location,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                else:
+                    if express_mode_api_key:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                    else:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            project=project,
+                            location=location,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
             except (ImportError, AttributeError):
                 from google.adk.sessions.vertex_ai_session_service_g3 import (
                     VertexAiSessionService,
                 )
 
-                self._tmpl_attrs["session_service"] = VertexAiSessionService(
-                    project=project,
-                    location=location,
-                    agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
-                )
+                if is_version_sufficient("1.18.0"):
+                    if express_mode_api_key:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            express_mode_api_key=express_mode_api_key,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                    else:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            project=project,
+                            location=location,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                else:
+                    if express_mode_api_key:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                    else:
+                        self._tmpl_attrs["session_service"] = VertexAiSessionService(
+                            project=project,
+                            location=location,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
 
         else:
             self._tmpl_attrs["session_service"] = InMemorySessionService()
@@ -780,11 +827,29 @@ def set_up(self):
                     VertexAiMemoryBankService,
                 )
 
-                self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
-                    project=project,
-                    location=location,
-                    agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
-                )
+                if is_version_sufficient("1.18.0"):
+                    if express_mode_api_key:
+                        self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
+                            express_mode_api_key=express_mode_api_key,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                    else:
+                        self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
+                            project=project,
+                            location=location,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                else:
+                    if express_mode_api_key:
+                        self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
+                    else:
+                        self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService(
+                            project=project,
+                            location=location,
+                            agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"),
+                        )
             except (ImportError, AttributeError):
                 # TODO(ysian): Handle this via _g3 import for google3.
                 pass