fix: Add absolutize_imports option when deploying to agent engine

PiperOrigin-RevId: 786749263

Author: yeesian

pyproject.toml

@@ -26,6 +26,7 @@ classifiers = [ # List of https://pypi.org/classifiers/
 dependencies = [
  # go/keep-sorted start
  "PyYAML>=6.0.2", # For APIHubToolset.
+ "absolufy-imports>=0.3.1", # For Agent Engine deployment.
  "anyio>=4.9.0;python_version>='3.10'", # For MCP Session Manager
  "authlib>=1.5.1", # For RestAPI Tool
  "click>=8.1.8", # For CLI tools

src/google/adk/cli/cli_deploy.py

@@ -59,7 +59,7 @@
 """
 
 _AGENT_ENGINE_APP_TEMPLATE = """
-from agent import root_agent
+from {app_name}.agent import root_agent
 from vertexai.preview.reasoning_engines import AdkApp
 
 adk_app = AdkApp(
@@ -254,6 +254,7 @@ def to_agent_engine(
  adk_app: str,
  staging_bucket: str,
  trace_to_cloud: bool,
+ absolutize_imports: bool = True,
  project: Optional[str] = None,
  region: Optional[str] = None,
  display_name: Optional[str] = None,
@@ -293,6 +294,8 @@ def to_agent_engine(
  region (str): Google Cloud region.
  staging_bucket (str): The GCS bucket for staging the deployment artifacts.
  trace_to_cloud (bool): Whether to enable Cloud Trace.
+ absolutize_imports (bool): Whether to absolutize imports. If True, all relative
+ imports will be converted to absolute import statements. Default is True.
  requirements_file (str): The filepath to the `requirements.txt` file to use.
  If not specified, the `requirements.txt` file in the `agent_folder` will
  be used.
@@ -301,14 +304,16 @@ def to_agent_engine(
  values of `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` will be
  overridden by `project` and `region` if they are specified.
  """
- # remove temp_folder if it exists
- if os.path.exists(temp_folder):
+ app_name = os.path.basename(agent_folder)
+ agent_src_path = os.path.join(temp_folder, app_name)
+ # remove agent_src_path if it exists
+ if os.path.exists(agent_src_path):
  click.echo('Removing existing files')
- shutil.rmtree(temp_folder)
+ shutil.rmtree(agent_src_path)
 
  try:
  click.echo('Copying agent source code...')
- shutil.copytree(agent_folder, temp_folder)
+ shutil.copytree(agent_folder, agent_src_path)
  click.echo('Copying agent source code complete.')
 
  click.echo('Initializing Vertex AI...')
@@ -317,13 +322,13 @@ def to_agent_engine(
  import vertexai
  from vertexai import agent_engines
 
- sys.path.append(temp_folder)
+ sys.path.append(temp_folder) # To register the adk_app operations
  project = _resolve_project(project)
 
  click.echo('Resolving files and dependencies...')
  if not requirements_file:
  # Attempt to read requirements from requirements.txt in the dir (if any).
- requirements_txt_path = os.path.join(temp_folder, 'requirements.txt')
+ requirements_txt_path = os.path.join(agent_src_path, 'requirements.txt')
  if not os.path.exists(requirements_txt_path):
  click.echo(f'Creating {requirements_txt_path}...')
  with open(requirements_txt_path, 'w', encoding='utf-8') as f:
@@ -333,7 +338,7 @@ def to_agent_engine(
  env_vars = None
  if not env_file:
  # Attempt to read the env variables from .env in the dir (if any).
- env_file = os.path.join(temp_folder, '.env')
+ env_file = os.path.join(agent_src_path, '.env')
  if os.path.exists(env_file):
  from dotenv import dotenv_values
 
@@ -371,21 +376,35 @@ def to_agent_engine(
  )
  click.echo('Vertex AI initialized.')
 
- adk_app_file = f'{adk_app}.py'
- with open(
- os.path.join(temp_folder, adk_app_file), 'w', encoding='utf-8'
- ) as f:
+ adk_app_file = os.path.join(temp_folder, f'{adk_app}.py')
+ with open(adk_app_file, 'w', encoding='utf-8') as f:
  f.write(
  _AGENT_ENGINE_APP_TEMPLATE.format(
- trace_to_cloud_option=trace_to_cloud
+ app_name=app_name,
+ trace_to_cloud_option=trace_to_cloud,
  )
  )
- click.echo(f'Created {os.path.join(temp_folder, adk_app_file)}')
+ click.echo(f'Created {adk_app_file}')
  click.echo('Files and dependencies resolved')
+ if absolutize_imports:
+ for root, _, files in os.walk(agent_src_path):
+ for file in files:
+ if file.endswith('.py'):
+ absolutize_imports_path = os.path.join(root, file)
+ try:
+ click.echo(
+ f'Running `absolufy-imports {absolutize_imports_path}`'
+ )
+ subprocess.run(
+ ['absolufy-imports', absolutize_imports_path],
+ cwd=temp_folder,
+ )
+ except Exception as e:
+ click.echo(f'The following exception was raised: {e}')
 
  click.echo('Deploying to agent engine...')
  agent_engine = agent_engines.ModuleAgent(
- module_name=adk_app,
+ module_name='agent_engine_app',
  agent_name='adk_app',
  register_operations={
  '': [

src/google/adk/cli/cli_tools_click.py

@@ -1055,6 +1055,16 @@ def cli_deploy_cloud_run(
  " any.)"
  ),
 )
+@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."
+ ),
+)
 @click.argument(
  "agent",
  type=click.Path(
@@ -1073,11 +1083,10 @@ def cli_deploy_agent_engine(
  temp_folder: str,
  env_file: str,
  requirements_file: str,
+ absolutize_imports: bool,
 ):
  """Deploys an agent to Agent Engine.
 
- AGENT: The path to the agent source code folder.
-
  Example:
 
  adk deploy agent_engine --project=[project] --region=[region]
@@ -1097,6 +1106,7 @@ def cli_deploy_agent_engine(
  temp_folder=temp_folder,
  env_file=env_file,
  requirements_file=requirements_file,
+ absolutize_imports=absolutize_imports,
  )
  except Exception as e:
  click.secho(f"Deploy failed: {e}", fg="red", err=True)