Support custom service account for agent engine via CLI

src/google/adk/cli/cli_deploy.py

@@ -83,6 +83,31 @@
 """
 
 
+def _resolve_param_from_env(
+    param_value: Optional[str],
+    env_var_name: str,
+    param_name: str,
+    env_vars: dict[str, str],
+    env_file: str,
+) -> Optional[str]:
+    """Resolves a parameter from environment variables, with CLI precedence."""
+    if env_var_name in env_vars:
+        env_value = env_vars.pop(env_var_name)
+        if env_value:
+            if param_value:
+                click.secho(
+                    f'Ignoring {env_var_name} in {env_file} as `--{param_name}` '
+                    'was explicitly passed and takes precedence.',
+                    fg='yellow',
+                )
+            else:
+                param_value = env_value
+                click.echo(
+                    f'{param_name}={param_value} set by {env_var_name} in {env_file}'
+                )
+    return param_value
+
+
 def _resolve_project(project_in_option: Optional[str]) -> str:
   if project_in_option:
     return project_in_option
@@ -350,6 +375,7 @@ def to_agent_engine(
     absolutize_imports: bool = True,
     project: Optional[str] = None,
     region: Optional[str] = None,
+    service_account: Optional[str] = None,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
     requirements_file: Optional[str] = None,
@@ -395,6 +421,10 @@ def to_agent_engine(
       import statements.
     project (str): Optional. Google Cloud project id.
     region (str): Optional. Google Cloud region.
+    service_account (str): Optional. Google Cloud service account to be used 
+      by the agent. It will override GOOGLE_CLOUD_SERVICE_ACCOUNT in the .env 
+      file (if it exists). If not provided, the AI Platform Reasoning Engine 
+      Service Agent will be used.
     display_name (str): Optional. The display name of the Agent Engine.
     description (str): Optional. The description of the Agent Engine.
     requirements_file (str): Optional. The filepath to the `requirements.txt`
@@ -494,30 +524,30 @@ def to_agent_engine(
 
       click.echo(f'Reading environment variables from {env_file}')
       env_vars = dotenv_values(env_file)
-      if 'GOOGLE_CLOUD_PROJECT' in env_vars:
-        env_project = env_vars.pop('GOOGLE_CLOUD_PROJECT')
-        if env_project:
-          if project:
-            click.secho(
-                'Ignoring GOOGLE_CLOUD_PROJECT in .env as `--project` was'
-                ' explicitly passed and takes precedence',
-                fg='yellow',
-            )
-          else:
-            project = env_project
-            click.echo(f'{project=} set by GOOGLE_CLOUD_PROJECT in {env_file}')
-      if 'GOOGLE_CLOUD_LOCATION' in env_vars:
-        env_region = env_vars.pop('GOOGLE_CLOUD_LOCATION')
-        if env_region:
-          if region:
-            click.secho(
-                'Ignoring GOOGLE_CLOUD_LOCATION in .env as `--region` was'
-                ' explicitly passed and takes precedence',
-                fg='yellow',
-            )
-          else:
-            region = env_region
-            click.echo(f'{region=} set by GOOGLE_CLOUD_LOCATION in {env_file}')
+      project = _resolve_param_from_env(
+        param_value=project,
+        env_var_name='GOOGLE_CLOUD_PROJECT',
+        param_name='project',
+        env_vars=env_vars,
+        env_file=env_file
+      )
+
+      region = _resolve_param_from_env(
+        param_value=region,
+        env_var_name='GOOGLE_CLOUD_LOCATION',
+        param_name='region',
+        env_vars=env_vars,
+        env_file=env_file
+      )
+
+      service_account = _resolve_param_from_env(
+        param_value=service_account,
+        env_var_name='GOOGLE_CLOUD_SERVICE_ACCOUNT',
+        param_name='service_account',
+        env_vars=env_vars,
+        env_file=env_file
+      )
+
     if env_vars:
       if 'env_vars' in agent_config:
         click.echo(
@@ -526,6 +556,9 @@ def to_agent_engine(
       agent_config['env_vars'] = env_vars
     # Set env_vars in agent_config to None if it is not set.
     agent_config['env_vars'] = agent_config.get('env_vars', env_vars)
+    # set service account if specified
+    if service_account:
+      agent_config['service_account'] = service_account
 
     vertexai.init(
         project=project,

src/google/adk/cli/cli_tools_click.py

@@ -1270,6 +1270,16 @@ def cli_deploy_cloud_run(
         " GOOGLE_CLOUD_LOCATION in the .env file (if it exists)."
     ),
 )
+@click.option(
+    "--service_account",
+    type=str,
+    help=(
+        "Optional. Google Cloud service account to be used by the agent."
+        "It will override GOOGLE_CLOUD_SERVICE_ACCOUNT in the .env file "
+        "(if it exists). If not provided, "
+        "the AI Platform Reasoning Engine Service Agent will be used."
+    ),
+)
 @click.option(
     "--staging_bucket",
     type=str,
@@ -1381,6 +1391,7 @@ def cli_deploy_agent_engine(
     agent: str,
     project: str,
     region: str,
+    service_account: Optional[str],
     staging_bucket: str,
     agent_engine_id: Optional[str],
     trace_to_cloud: bool,
@@ -1406,6 +1417,7 @@ def cli_deploy_agent_engine(
         agent_folder=agent,
         project=project,
         region=region,
+        service_account=service_account,
         staging_bucket=staging_bucket,
         agent_engine_id=agent_engine_id,
         trace_to_cloud=trace_to_cloud,