From 2409c3ef192262c80f5328121f6dc4f34265f5cf Mon Sep 17 00:00:00 2001 From: Yeesian Ng Date: Mon, 2 Jun 2025 08:18:00 -0700 Subject: [PATCH] feat: Add agent engine as a deployment option to the ADK CLI PiperOrigin-RevId: 766199746 --- pyproject.toml | 36 +++--- src/google/adk/cli/cli_deploy.py | 151 ++++++++++++++++++++++++++ src/google/adk/cli/cli_tools_click.py | 123 +++++++++++++++++++++ 3 files changed, 292 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a97c5ef..1c66e92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,26 +25,26 @@ classifiers = [ # List of https://pypi.org/classifiers/ ] dependencies = [ # go/keep-sorted start - "authlib>=1.5.1", # For RestAPI Tool - "click>=8.1.8", # For CLI tools - "fastapi>=0.115.0", # FastAPI framework - "google-api-python-client>=2.157.0", # Google API client discovery - "google-cloud-aiplatform>=1.87.0", # For VertexAI integrations, e.g. example store. - "google-cloud-secret-manager>=2.22.0", # Fetching secrets in RestAPI Tool - "google-cloud-speech>=2.30.0", # For Audio Transcription - "google-cloud-storage>=2.18.0, <3.0.0", # For GCS Artifact service - "google-genai>=1.17.0", # Google GenAI SDK - "graphviz>=0.20.2", # Graphviz for graph rendering - "mcp>=1.8.0;python_version>='3.10'", # For MCP Toolset - "opentelemetry-api>=1.31.0", # OpenTelemetry + "authlib>=1.5.1", # For RestAPI Tool + "click>=8.1.8", # For CLI tools + "fastapi>=0.115.0", # FastAPI framework + "google-api-python-client>=2.157.0", # Google API client discovery + "google-cloud-aiplatform[agent_engines]>=1.95.1", # For VertexAI integrations, e.g. example store. + "google-cloud-secret-manager>=2.22.0", # Fetching secrets in RestAPI Tool + "google-cloud-speech>=2.30.0", # For Audio Transcription + "google-cloud-storage>=2.18.0, <3.0.0", # For GCS Artifact service + "google-genai>=1.17.0", # Google GenAI SDK + "graphviz>=0.20.2", # Graphviz for graph rendering + "mcp>=1.8.0;python_version>='3.10'", # For MCP Toolset + "opentelemetry-api>=1.31.0", # OpenTelemetry "opentelemetry-exporter-gcp-trace>=1.9.0", "opentelemetry-sdk>=1.31.0", - "pydantic>=2.0, <3.0.0", # For data validation/models - "python-dotenv>=1.0.0", # To manage environment variables - "PyYAML>=6.0.2", # For APIHubToolset. - "sqlalchemy>=2.0", # SQL database ORM - "tzlocal>=5.3", # Time zone utilities - "uvicorn>=0.34.0", # ASGI server for FastAPI + "pydantic>=2.0, <3.0.0", # For data validation/models + "python-dotenv>=1.0.0", # To manage environment variables + "PyYAML>=6.0.2", # For APIHubToolset. + "sqlalchemy>=2.0", # SQL database ORM + "tzlocal>=5.3", # Time zone utilities + "uvicorn>=0.34.0", # ASGI server for FastAPI # go/keep-sorted end ] dynamic = ["version"] diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index a478799..8ff2755 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -58,6 +58,16 @@ EXPOSE {port} CMD adk {command} --port={port} {host_option} {session_db_option} {trace_to_cloud_option} "/app/agents" """ +_AGENT_ENGINE_APP_TEMPLATE = """ +from agent import root_agent +from vertexai.preview.reasoning_engines import AdkApp + +adk_app = AdkApp( + agent=root_agent, + enable_tracing={trace_to_cloud_option}, +) +""" + def _resolve_project(project_in_option: Optional[str]) -> str: if project_in_option: @@ -197,3 +207,144 @@ def to_cloud_run( finally: click.echo(f'Cleaning up the temp folder: {temp_folder}') shutil.rmtree(temp_folder) + + +def to_agent_engine( + *, + agent_folder: str, + temp_folder: str, + adk_app: str, + project: str, + region: str, + staging_bucket: str, + trace_to_cloud: bool, + requirements_file: Optional[str] = None, + env_file: Optional[str] = None, +): + """Deploys an agent to Vertex AI Agent Engine. + + `agent_folder` should contain the following files: + + - __init__.py + - agent.py + - .py (optional, for customization; will be autogenerated otherwise) + - requirements.txt (optional, for additional dependencies) + - .env (optional, for environment variables) + - ... (other required source files) + + The contents of `adk_app` should look something like: + + ``` + from agent import root_agent + from vertexai.preview.reasoning_engines import AdkApp + + adk_app = AdkApp( + agent=root_agent, + enable_tracing=True, + ) + ``` + + Args: + agent_folder (str): The folder (absolute path) containing the agent source + code. + temp_folder (str): The temp folder for the generated Agent Engine source + files. It will be replaced with the generated files if it already exists. + project (str): Google Cloud project id. + 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. + 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. + env_file (str): The filepath to the `.env` file for environment variables. + If not specified, the `.env` file in the `agent_folder` will be used. + """ + # remove temp_folder if it exists + if os.path.exists(temp_folder): + click.echo('Removing existing files') + shutil.rmtree(temp_folder) + + try: + click.echo('Copying agent source code...') + shutil.copytree(agent_folder, temp_folder) + click.echo('Copying agent source code complete.') + + click.echo('Initializing Vertex AI...') + import sys + + import vertexai + from vertexai import agent_engines + + sys.path.append(temp_folder) + + vertexai.init( + project=_resolve_project(project), + location=region, + staging_bucket=staging_bucket, + ) + click.echo('Vertex AI initialized.') + + 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') + 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: + f.write('google-cloud-aiplatform[adk,agent_engines]') + click.echo(f'Created {requirements_txt_path}') + requirements_file = requirements_txt_path + 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') + if os.path.exists(env_file): + from dotenv import dotenv_values + + click.echo(f'Reading environment variables from {env_file}') + env_vars = dotenv_values(env_file) + + adk_app_file = f'{adk_app}.py' + with open( + os.path.join(temp_folder, adk_app_file), 'w', encoding='utf-8' + ) as f: + f.write( + _AGENT_ENGINE_APP_TEMPLATE.format( + trace_to_cloud_option=trace_to_cloud + ) + ) + click.echo(f'Created {os.path.join(temp_folder, adk_app_file)}') + click.echo('Files and dependencies resolved') + + click.echo('Deploying to agent engine...') + agent_engine = agent_engines.ModuleAgent( + module_name=adk_app, + agent_name='adk_app', + register_operations={ + '': [ + 'get_session', + 'list_sessions', + 'create_session', + 'delete_session', + ], + 'async': [ + 'async_get_session', + 'async_list_sessions', + 'async_create_session', + 'async_delete_session', + ], + 'async_stream': ['async_stream_query'], + 'stream': ['stream_query', 'streaming_agent_run_with_events'], + }, + sys_paths=[temp_folder[1:]], + ) + + agent_engines.create( + agent_engine=agent_engine, + requirements=requirements_file, + env_vars=env_vars, + extra_packages=[temp_folder], + ) + finally: + click.echo(f'Cleaning up the temp folder: {temp_folder}') + shutil.rmtree(temp_folder) diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 3d2e5d3..e8da225 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -767,3 +767,126 @@ def cli_deploy_cloud_run( ) except Exception as e: click.secho(f"Deploy failed: {e}", fg="red", err=True) + + +@deploy.command("agent_engine") +@click.option( + "--project", + type=str, + help="Required. Google Cloud project to deploy the agent.", +) +@click.option( + "--region", + type=str, + help="Required. Google Cloud region to deploy the agent.", +) +@click.option( + "--staging_bucket", + type=str, + help="Required. GCS bucket for staging the deployment artifacts.", +) +@click.option( + "--trace_to_cloud", + type=bool, + is_flag=True, + show_default=True, + default=False, + help="Optional. Whether to enable Cloud Trace for Agent Engine.", +) +@click.option( + "--adk_app", + type=str, + default="agent_engine_app", + help=( + "Optional. Python file for defining the ADK application" + " (default: a file named agent_engine_app.py)" + ), +) +@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"), + ), + 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)." + ), +) +@click.option( + "--env_file", + type=str, + default="", + help=( + "Optional. The filepath to the `.env` file for environment variables." + " (default: the `.env` file in the `agent` directory, if any.)" + ), +) +@click.option( + "--requirements_file", + type=str, + default="", + help=( + "Optional. The filepath to the `requirements.txt` file to use." + " (default: the `requirements.txt` file in the `agent` directory, if" + " any.)" + ), +) +@click.argument( + "agent", + type=click.Path( + exists=True, dir_okay=True, file_okay=False, resolve_path=True + ), +) +def cli_deploy_agent_engine( + agent: str, + project: str, + region: str, + staging_bucket: str, + trace_to_cloud: bool, + adk_app: str, + temp_folder: str, + env_file: str, + requirements_file: str, +): + """Deploys an agent to Agent Engine. + + Args: + agent (str): Required. The path to the agent to be deloyed. + project (str): Required. Google Cloud project to deploy the agent. + region (str): Required. Google Cloud region to deploy the agent. + staging_bucket (str): Required. GCS bucket for staging the deployment + artifacts. + trace_to_cloud (bool): Required. Whether to enable Cloud Trace. + adk_app (str): Required. Python file for defining the ADK application. + temp_folder (str): Required. The folder for the generated Agent Engine + files. If the folder already exists, its contents will be replaced. + env_file (str): Required. The filepath to the `.env` file for environment + variables. If it is an empty string, the `.env` file in the `agent` + directory will be used if it exists. + requirements_file (str): Required. The filepath to the `requirements.txt` + file to use. If it is an empty string, the `requirements.txt` file in the + `agent` directory will be used if exists. + + Example: + + adk deploy agent_engine --project=[project] --region=[region] + --staging_bucket=[staging_bucket] path/to/my_agent + """ + try: + cli_deploy.to_agent_engine( + agent_folder=agent, + project=project, + region=region, + staging_bucket=staging_bucket, + trace_to_cloud=trace_to_cloud, + adk_app=adk_app, + temp_folder=temp_folder, + env_file=env_file, + requirements_file=requirements_file, + ) + except Exception as e: + click.secho(f"Deploy failed: {e}", fg="red", err=True)