feat: Add agent engine as a deployment option to the ADK CLI

PiperOrigin-RevId: 766199746
This commit is contained in:
Yeesian Ng 2025-06-02 08:18:00 -07:00 committed by Copybara-Service
parent 8d36dbda52
commit 2409c3ef19
3 changed files with 292 additions and 18 deletions

View File

@ -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"]

View File

@ -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
- <adk_app>.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)

View File

@ -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)