feat: Add support for Vertex AI Express Mode when deploying to Agent Engine

Co-authored-by: Yeesian Ng <ysian@google.com>
PiperOrigin-RevId: 828178479

pyproject.toml

@@ -26,13 +26,12 @@ classifiers = [ # List of https://pypi.org/classifiers/
 dependencies = [
   # go/keep-sorted start
   "PyYAML>=6.0.2, <7.0.0",                                   # For APIHubToolset.
-  "absolufy-imports>=0.3.1, <1.0.0",                         # For Agent Engine deployment.
   "anyio>=4.9.0, <5.0.0;python_version>='3.10'",             # For MCP Session Manager
   "authlib>=1.5.1, <2.0.0",                                  # For RestAPI Tool
   "click>=8.1.8, <9.0.0",                                    # For CLI tools
   "fastapi>=0.115.0, <1.119.0",                              # FastAPI framework
   "google-api-python-client>=2.157.0, <3.0.0",               # Google API client discovery
-  "google-cloud-aiplatform[agent_engines]>=1.121.0, <2.0.0", # For VertexAI integrations, e.g. example store.
+  "google-cloud-aiplatform[agent_engines] @ git+https://github.com/googleapis/python-aiplatform.git@bf1851e59cb34e63b509a2a610e72691e1c4ca28", # For VertexAI integrations, e.g. example store.
   "google-cloud-bigtable>=2.32.0",                           # For Bigtable database
   "google-cloud-discoveryengine>=0.13.12, <0.14.0",          # For Discovery Engine Search Tool
   "google-cloud-secret-manager>=2.22.0, <3.0.0",             # Fetching secrets in RestAPI Tool

src/google/adk/cli/cli_tools_click.py

@@ -1462,26 +1462,46 @@ def cli_deploy_cloud_run(
 
 
 @deploy.command("agent_engine")
+@click.option(
+    "--api_key",
+    type=str,
+    default=None,
+    help=(
+        "Optional. The API key to use for Express Mode. If not"
+        " provided, the API key from the GOOGLE_API_KEY environment variable"
+        " will be used. It will only be used if GOOGLE_GENAI_USE_VERTEXAI is"
+        " true. (It will override GOOGLE_API_KEY in the .env file if it"
+        " exists.)"
+    ),
+)
 @click.option(
     "--project",
     type=str,
+    default=None,
     help=(
-        "Required. Google Cloud project to deploy the agent. It will override"
-        " GOOGLE_CLOUD_PROJECT in the .env file (if it exists)."
+        "Optional. Google Cloud project to deploy the agent. It will override"
+        " GOOGLE_CLOUD_PROJECT in the .env file (if it exists). It will be"
+        " ignored if api_key is set."
     ),
 )
 @click.option(
     "--region",
     type=str,
+    default=None,
     help=(
-        "Required. Google Cloud region to deploy the agent. It will override"
-        " GOOGLE_CLOUD_LOCATION in the .env file (if it exists)."
+        "Optional. Google Cloud region to deploy the agent. It will override"
+        " GOOGLE_CLOUD_LOCATION in the .env file (if it exists). It will be"
+        " ignored if api_key is set."
     ),
 )
 @click.option(
     "--staging_bucket",
     type=str,
-    help="Required. GCS bucket for staging the deployment artifacts.",
+    default=None,
+    help=(
+        "Optional. GCS bucket for staging the deployment artifacts. It will be"
+        " ignored if api_key is set."
+    ),
 )
 @click.option(
     "--agent_engine_id",
@@ -1489,9 +1509,12 @@ def cli_deploy_cloud_run(
     default=None,
     help=(
         "Optional. ID of the Agent Engine instance to update if it exists"
-        " (default: None, which means a new instance will be created)."
-        " The corresponding resource name in Agent Engine will be:"
+        " (default: None, which means a new instance will be created). If"
+        " project and region are set, this should be the resource ID, and the"
+        " corresponding resource name in Agent Engine will be:"
         " `projects/{project}/locations/{region}/reasoningEngines/{agent_engine_id}`."
+        " If api_key is set, then agent_engine_id is required to be the full"
+        " resource name (i.e. `projects/*/locations/*/reasoningEngines/*`)."
     ),
 )
 @click.option(
@@ -1528,15 +1551,20 @@ def cli_deploy_cloud_run(
 @click.option(
     "--temp_folder",
     type=str,
-    default=os.path.join(
-        tempfile.gettempdir(),
-        "agent_engine_deploy_src",
-        datetime.now().strftime("%Y%m%d_%H%M%S"),
-    ),
+    default=None,
     help=(
         "Optional. Temp folder for the generated Agent Engine source files."
         " If the folder already exists, its contents will be removed."
-        " (default: a timestamped folder in the system temp directory)."
+        " (default: a timestamped folder in the current working directory)."
+    ),
+)
+@click.option(
+    "--adk_app_object",
+    type=str,
+    default=None,
+    help=(
+        "Optional. Python object corresponding to the root ADK agent or app."
+        " It can only be `root_agent` or `app`. (default: `root_agent`)"
     ),
 )
 @click.option(
@@ -1561,12 +1589,8 @@ def cli_deploy_cloud_run(
 @click.option(
     "--absolutize_imports",
     type=bool,
-    default=True,
-    help=(
-        "Optional. Whether to absolutize imports. If True, all relative imports"
-        " will be converted to absolute import statements (default: True)."
-        " NOTE: This flag is temporary and will be removed in the future."
-    ),
+    default=False,
+    help=" NOTE: This flag is deprecated and will be removed in the future.",
 )
 @click.option(
     "--agent_engine_config_file",
@@ -1587,15 +1611,17 @@ def cli_deploy_cloud_run(
 )
 def cli_deploy_agent_engine(
     agent: str,
-    project: str,
-    region: str,
-    staging_bucket: str,
+    project: Optional[str],
+    region: Optional[str],
+    staging_bucket: Optional[str],
     agent_engine_id: Optional[str],
     trace_to_cloud: Optional[bool],
+    api_key: Optional[str],
     display_name: str,
     description: str,
     adk_app: str,
-    temp_folder: str,
+    adk_app_object: Optional[str],
+    temp_folder: Optional[str],
     env_file: str,
     requirements_file: str,
     absolutize_imports: bool,
@@ -1605,9 +1631,13 @@ def cli_deploy_agent_engine(
 
   Example:
 
+    # With Express Mode API Key
+    adk deploy agent_engine --api_key=[api_key] my_agent
+
+    # With Google Cloud Project and Region
     adk deploy agent_engine --project=[project] --region=[region]
       --staging_bucket=[staging_bucket] --display_name=[app_name]
-      path/to/my_agent
+      my_agent
   """
   try:
     cli_deploy.to_agent_engine(
@@ -1617,6 +1647,8 @@ def cli_deploy_agent_engine(
         staging_bucket=staging_bucket,
         agent_engine_id=agent_engine_id,
         trace_to_cloud=trace_to_cloud,
+        api_key=api_key,
+        adk_app_object=adk_app_object,
         display_name=display_name,
         description=description,
         adk_app=adk_app,

tests/unittests/cli/utils/test_cli_deploy.py

@@ -95,34 +95,6 @@ def agent_dir(tmp_path: Path) -> Callable[[bool, bool], Path]:
   return _factory
 
 
-@pytest.fixture
-def mock_vertex_ai(
-    monkeypatch: pytest.MonkeyPatch,
-) -> Generator[mock.MagicMock, None, None]:
-  """Mocks the entire vertexai module and its sub-modules."""
-  mock_vertexai = mock.MagicMock()
-  mock_agent_engines = mock.MagicMock()
-  mock_vertexai.agent_engines = mock_agent_engines
-  mock_vertexai.init = mock.MagicMock()
-  mock_agent_engines.create = mock.MagicMock()
-  mock_agent_engines.ModuleAgent = mock.MagicMock(
-      return_value="mock-agent-engine-object"
-  )
-
-  sys.modules["vertexai"] = mock_vertexai
-  sys.modules["vertexai.agent_engines"] = mock_agent_engines
-
-  mock_dotenv = mock.MagicMock()
-  mock_dotenv.dotenv_values = mock.MagicMock(return_value={"FILE_VAR": "value"})
-  sys.modules["dotenv"] = mock_dotenv
-
-  yield mock_vertexai
-
-  del sys.modules["vertexai"]
-  del sys.modules["vertexai.agent_engines"]
-  del sys.modules["dotenv"]
-
-
 # _resolve_project
 def test_resolve_project_with_option() -> None:
   """It should return the explicit project value untouched."""
@@ -216,80 +188,6 @@ def test_get_service_option_by_adk_version(
   assert actual.rstrip() == expected.rstrip()
 
 
-@pytest.mark.usefixtures("mock_vertex_ai")
-@pytest.mark.parametrize("has_reqs", [True, False])
-@pytest.mark.parametrize("has_env", [True, False])
-def test_to_agent_engine_happy_path(
-    monkeypatch: pytest.MonkeyPatch,
-    agent_dir: Callable[[bool, bool], Path],
-    tmp_path: Path,
-    has_reqs: bool,
-    has_env: bool,
-) -> None:
-  """
-  Tests the happy path for the `to_agent_engine` function.
-  """
-  src_dir = agent_dir(has_reqs, has_env)
-  temp_folder = tmp_path / "build"
-  app_name = src_dir.name
-  rmtree_recorder = _Recorder()
-
-  monkeypatch.setattr(shutil, "rmtree", rmtree_recorder)
-
-  cli_deploy.to_agent_engine(
-      agent_folder=str(src_dir),
-      temp_folder=str(temp_folder),
-      adk_app="my_adk_app",
-      staging_bucket="gs://my-staging-bucket",
-      trace_to_cloud=True,
-      project="my-gcp-project",
-      region="us-central1",
-      display_name="My Test Agent",
-      description="A test agent.",
-  )
-
-  assert (temp_folder / app_name / "agent.py").is_file()
-  assert (temp_folder / app_name / "__init__.py").is_file()
-
-  adk_app_path = temp_folder / "my_adk_app.py"
-  assert adk_app_path.is_file()
-  content = adk_app_path.read_text()
-  assert f"from {app_name}.agent import root_agent" in content
-  assert "adk_app = AdkApp(" in content
-  assert "enable_tracing=True" in content
-
-  reqs_path = temp_folder / app_name / "requirements.txt"
-  assert reqs_path.is_file()
-  if not has_reqs:
-    assert "google-cloud-aiplatform[adk,agent_engines]" in reqs_path.read_text()
-
-  vertexai = sys.modules["vertexai"]
-  vertexai.init.assert_called_once_with(
-      project="my-gcp-project",
-      location="us-central1",
-      staging_bucket="gs://my-staging-bucket",
-  )
-
-  dotenv = sys.modules["dotenv"]
-  if has_env:
-    dotenv.dotenv_values.assert_called_once()
-    expected_env_vars = {"FILE_VAR": "value"}
-  else:
-    dotenv.dotenv_values.assert_not_called()
-    expected_env_vars = None
-
-  vertexai.agent_engines.create.assert_called_once()
-  create_kwargs = vertexai.agent_engines.create.call_args.kwargs
-  assert create_kwargs["agent_engine"] == "mock-agent-engine-object"
-  assert create_kwargs["display_name"] == "My Test Agent"
-  assert create_kwargs["description"] == "A test agent."
-  assert create_kwargs["requirements"] == str(reqs_path)
-  assert create_kwargs["extra_packages"] == [str(temp_folder)]
-  assert create_kwargs["env_vars"] == expected_env_vars
-
-  assert str(rmtree_recorder.get_last_call_args()[0]) == str(temp_folder)
-
-
 @pytest.mark.parametrize("include_requirements", [True, False])
 def test_to_gke_happy_path(
     monkeypatch: pytest.MonkeyPatch,