From d4b2a8b49f98a9991cb44ac7ec6e538b81a08664 Mon Sep 17 00:00:00 2001 From: Yeesian Ng Date: Tue, 4 Nov 2025 16:23:36 -0800 Subject: [PATCH] feat: Add support for Vertex AI Express Mode when deploying to Agent Engine Co-authored-by: Yeesian Ng PiperOrigin-RevId: 828178479 --- pyproject.toml | 3 +- src/google/adk/cli/cli_deploy.py | 458 ++++++++++++++++--- src/google/adk/cli/cli_tools_click.py | 80 +++- tests/unittests/cli/utils/test_cli_deploy.py | 102 ----- 4 files changed, 444 insertions(+), 199 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0ba44779..149a6ec2 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index 1fc79204..335c786b 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +from datetime import datetime import json import os import shutil @@ -63,7 +64,7 @@ CMD adk {command} --port={port} {host_option} {service_option} {trace_to_cloud_o """ _AGENT_ENGINE_APP_TEMPLATE: Final[str] = """ -from vertexai.preview.reasoning_engines import AdkApp +from vertexai.agent_engines import AdkApp if {is_config_agent}: from google.adk.agents import config_agent_utils @@ -74,14 +75,298 @@ if {is_config_agent}: # This path is used to support the file structure in Agent Engine. root_agent = config_agent_utils.from_config("./{temp_folder}/{app_name}/root_agent.yaml") else: - from {app_name}.agent import root_agent + from .agent import {adk_app_object} + +if {express_mode}: # Whether or not to use Express Mode + import os + import vertexai + vertexai.init(api_key=os.environ.get("GOOGLE_API_KEY")) adk_app = AdkApp( - agent=root_agent, - enable_tracing={trace_to_cloud_option}, + {adk_app_type}={adk_app_object}, + enable_tracing={trace_to_cloud_option}, ) """ +_AGENT_ENGINE_CLASS_METHODS = [ + { + 'name': 'get_session', + 'description': ( + 'Deprecated. Use async_get_session instead.\n\n Get a' + ' session for the given user.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string'}, + }, + 'required': ['user_id', 'session_id'], + 'type': 'object', + }, + 'api_mode': '', + }, + { + 'name': 'list_sessions', + 'description': ( + 'Deprecated. Use async_list_sessions instead.\n\n List' + ' sessions for the given user.\n ' + ), + 'parameters': { + 'properties': {'user_id': {'type': 'string'}}, + 'required': ['user_id'], + 'type': 'object', + }, + 'api_mode': '', + }, + { + 'name': 'create_session', + 'description': ( + 'Deprecated. Use async_create_session instead.\n\n Creates a' + ' new session.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string', 'nullable': True}, + 'state': {'type': 'object', 'nullable': True}, + }, + 'required': ['user_id'], + 'type': 'object', + }, + 'api_mode': '', + }, + { + 'name': 'delete_session', + 'description': ( + 'Deprecated. Use async_delete_session instead.\n\n Deletes a' + ' session for the given user.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string'}, + }, + 'required': ['user_id', 'session_id'], + 'type': 'object', + }, + 'api_mode': '', + }, + { + 'name': 'async_get_session', + 'description': ( + 'Get a session for the given user.\n\n Args:\n ' + ' user_id (str):\n Required. The ID of the user.\n ' + ' session_id (str):\n Required. The ID of' + ' the session.\n **kwargs (dict[str, Any]):\n ' + ' Optional. Additional keyword arguments to pass to the\n ' + ' session service.\n\n Returns:\n ' + ' Session: The session instance (if any). It returns None if the\n ' + ' session is not found.\n\n Raises:\n ' + ' RuntimeError: If the session is not found.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string'}, + }, + 'required': ['user_id', 'session_id'], + 'type': 'object', + }, + 'api_mode': 'async', + }, + { + 'name': 'async_list_sessions', + 'description': ( + 'List sessions for the given user.\n\n Args:\n ' + ' user_id (str):\n Required. The ID of the user.\n ' + ' **kwargs (dict[str, Any]):\n Optional.' + ' Additional keyword arguments to pass to the\n ' + ' session service.\n\n Returns:\n ' + ' ListSessionsResponse: The list of sessions.\n ' + ), + 'parameters': { + 'properties': {'user_id': {'type': 'string'}}, + 'required': ['user_id'], + 'type': 'object', + }, + 'api_mode': 'async', + }, + { + 'name': 'async_create_session', + 'description': ( + 'Creates a new session.\n\n Args:\n user_id' + ' (str):\n Required. The ID of the user.\n ' + ' session_id (str):\n Optional. The ID of the' + ' session. If not provided, an ID\n will be be' + ' generated for the session.\n state (dict[str, Any]):\n' + ' Optional. The initial state of the session.\n ' + ' **kwargs (dict[str, Any]):\n Optional.' + ' Additional keyword arguments to pass to the\n ' + ' session service.\n\n Returns:\n Session: The' + ' newly created session instance.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string', 'nullable': True}, + 'state': {'type': 'object', 'nullable': True}, + }, + 'required': ['user_id'], + 'type': 'object', + }, + 'api_mode': 'async', + }, + { + 'name': 'async_delete_session', + 'description': ( + 'Deletes a session for the given user.\n\n Args:\n ' + ' user_id (str):\n Required. The ID of the user.\n ' + ' session_id (str):\n Required. The ID of' + ' the session.\n **kwargs (dict[str, Any]):\n ' + ' Optional. Additional keyword arguments to pass to the\n ' + ' session service.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string'}, + }, + 'required': ['user_id', 'session_id'], + 'type': 'object', + }, + 'api_mode': 'async', + }, + { + 'name': 'async_add_session_to_memory', + 'description': ( + 'Generates memories.\n\n Args:\n session' + ' (Dict[str, Any]):\n Required. The session to use' + ' for generating memories. It should\n be a' + ' dictionary representing an ADK Session object, e.g.\n ' + ' session.model_dump(mode="json").\n ' + ), + 'parameters': { + 'properties': { + 'session': {'additionalProperties': True, 'type': 'object'} + }, + 'required': ['session'], + 'type': 'object', + }, + 'api_mode': 'async', + }, + { + 'name': 'async_search_memory', + 'description': ( + 'Searches memories for the given user.\n\n Args:\n ' + ' user_id: The id of the user.\n query: The query to' + ' match the memories on.\n\n Returns:\n A' + ' SearchMemoryResponse containing the matching memories.\n ' + ), + 'parameters': { + 'properties': { + 'user_id': {'type': 'string'}, + 'query': {'type': 'string'}, + }, + 'required': ['user_id', 'query'], + 'type': 'object', + }, + 'api_mode': 'async', + }, + { + 'name': 'stream_query', + 'description': ( + 'Deprecated. Use async_stream_query instead.\n\n Streams' + ' responses from the ADK application in response to a message.\n\n ' + ' Args:\n message (Union[str, Dict[str, Any]]):\n ' + ' Required. The message to stream responses for.\n ' + ' user_id (str):\n Required. The ID of the' + ' user.\n session_id (str):\n Optional.' + ' The ID of the session. If not provided, a new\n ' + ' session will be created for the user.\n run_config' + ' (Optional[Dict[str, Any]]):\n Optional. The run' + ' config to use for the query. If you want to\n pass' + ' in a `run_config` pydantic object, you can pass in a dict\n ' + ' representing it as' + ' `run_config.model_dump(mode="json")`.\n **kwargs' + ' (dict[str, Any]):\n Optional. Additional keyword' + ' arguments to pass to the\n runner.\n\n ' + ' Yields:\n The output of querying the ADK' + ' application.\n ' + ), + 'parameters': { + 'properties': { + 'message': { + 'anyOf': [ + {'type': 'string'}, + {'additionalProperties': True, 'type': 'object'}, + ] + }, + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string', 'nullable': True}, + 'run_config': {'type': 'object', 'nullable': True}, + }, + 'required': ['message', 'user_id'], + 'type': 'object', + }, + 'api_mode': 'stream', + }, + { + 'name': 'async_stream_query', + 'description': ( + 'Streams responses asynchronously from the ADK application.\n\n ' + ' Args:\n message (str):\n Required.' + ' The message to stream responses for.\n user_id' + ' (str):\n Required. The ID of the user.\n ' + ' session_id (str):\n Optional. The ID of the' + ' session. If not provided, a new\n session will be' + ' created for the user.\n run_config (Optional[Dict[str,' + ' Any]]):\n Optional. The run config to use for the' + ' query. If you want to\n pass in a `run_config`' + ' pydantic object, you can pass in a dict\n ' + ' representing it as `run_config.model_dump(mode="json")`.\n ' + ' **kwargs (dict[str, Any]):\n Optional.' + ' Additional keyword arguments to pass to the\n ' + ' runner.\n\n Yields:\n Event dictionaries' + ' asynchronously.\n ' + ), + 'parameters': { + 'properties': { + 'message': { + 'anyOf': [ + {'type': 'string'}, + {'additionalProperties': True, 'type': 'object'}, + ] + }, + 'user_id': {'type': 'string'}, + 'session_id': {'type': 'string', 'nullable': True}, + 'run_config': {'type': 'object', 'nullable': True}, + }, + 'required': ['message', 'user_id'], + 'type': 'object', + }, + 'api_mode': 'async_stream', + }, + { + 'name': 'streaming_agent_run_with_events', + 'description': ( + 'Streams responses asynchronously from the ADK application.\n\n ' + ' In general, you should use `async_stream_query` instead, as it' + ' has a\n more structured API and works with the respective' + ' ADK services that\n you have defined for the AdkApp. This' + ' method is primarily meant for\n invocation from' + ' AgentSpace.\n\n Args:\n request_json (str):\n ' + ' Required. The request to stream responses for.\n ' + ' ' + ), + 'parameters': { + 'properties': {'request_json': {'type': 'string'}}, + 'required': ['request_json'], + 'type': 'object', + }, + 'api_mode': 'async_stream', + }, +] + def _resolve_project(project_in_option: Optional[str]) -> str: if project_in_option: @@ -342,10 +627,12 @@ def to_cloud_run( def to_agent_engine( *, agent_folder: str, - temp_folder: str, + temp_folder: Optional[str] = None, adk_app: str, staging_bucket: str, trace_to_cloud: Optional[bool] = None, + api_key: Optional[str] = None, + adk_app_object: Optional[str] = None, agent_engine_id: Optional[str] = None, absolutize_imports: bool = True, project: Optional[str] = None, @@ -370,12 +657,11 @@ def to_agent_engine( The contents of `adk_app` should look something like: ``` - from agent import root_agent - from vertexai.preview.reasoning_engines import AdkApp + from agent import + from vertexai.agent_engines import AdkApp adk_app = AdkApp( - agent=root_agent, - enable_tracing=True, + agent=, # or `app=` ) ``` @@ -388,6 +674,11 @@ def to_agent_engine( instance. staging_bucket (str): The GCS bucket for staging the deployment artifacts. trace_to_cloud (bool): Whether to enable Cloud Trace. + api_key (str): 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. + adk_app_object (str): Optional. The Python object corresponding to the root + ADK agent or app. Defaults to `root_agent` if not specified. agent_engine_id (str): Optional. The ID of the Agent Engine instance to update. If not specified, a new Agent Engine instance will be created. absolutize_imports (bool): Optional. Default is True. Whether to absolutize @@ -409,7 +700,22 @@ def to_agent_engine( `agent_folder` will be used. """ app_name = os.path.basename(agent_folder) - agent_src_path = os.path.join(temp_folder, app_name) + display_name = display_name or app_name + parent_folder = os.path.dirname(agent_folder) + if parent_folder != os.getcwd(): + click.echo(f'Please deploy from the project dir: {parent_folder}') + return + tmp_app_name = app_name + '_tmp' + datetime.now().strftime('%Y%m%d_%H%M%S') + temp_folder = temp_folder or tmp_app_name + agent_src_path = os.path.join(parent_folder, temp_folder) + click.echo(f'Staging all files in: {agent_src_path}') + adk_app_object = adk_app_object or 'root_agent' + if adk_app_object not in ['root_agent', 'app']: + click.echo( + f'Invalid adk_app_object: {adk_app_object}. Please use "root_agent"' + ' or "app".' + ) + return # remove agent_src_path if it exists if os.path.exists(agent_src_path): click.echo('Removing existing files') @@ -427,17 +733,12 @@ def to_agent_engine( shutil.copytree(agent_folder, agent_src_path, ignore=ignore_patterns) 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) # To register the adk_app operations project = _resolve_project(project) click.echo('Resolving files and dependencies...') agent_config = {} + if staging_bucket: + agent_config['staging_bucket'] = staging_bucket if not agent_engine_config_file: # Attempt to read the agent engine config from .agent_engine_config.json in the dir (if any). agent_engine_config_file = os.path.join( @@ -460,10 +761,6 @@ def to_agent_engine( f'Overriding description in agent engine config with {description}' ) agent_config['description'] = description - if agent_config.get('extra_packages'): - agent_config['extra_packages'].append(temp_folder) - else: - agent_config['extra_packages'] = [temp_folder] if not requirements_file: # Attempt to read requirements from requirements.txt in the dir (if any). @@ -471,21 +768,26 @@ def to_agent_engine( 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]') + f.write( + 'google-cloud-aiplatform[adk,agent_engines] @ ' + 'git+https://github.com/googleapis/python-aiplatform.git@' + 'bf1851e59cb34e63b509a2a610e72691e1c4ca28' + ) click.echo(f'Created {requirements_txt_path}') - agent_config['requirements'] = agent_config.get( + agent_config['requirements_file'] = agent_config.get( 'requirements', requirements_txt_path, ) else: - if 'requirements' in agent_config: + if 'requirements_file' in agent_config: click.echo( 'Overriding requirements in agent engine config with ' f'{requirements_file}' ) - agent_config['requirements'] = requirements_file + agent_config['requirements_file'] = requirements_file + agent_config['requirements_file'] = f'{temp_folder}/requirements.txt' - env_vars = None + env_vars = {} if not env_file: # Attempt to read the env variables from .env in the dir (if any). env_file = os.path.join(agent_folder, '.env') @@ -518,6 +820,20 @@ def to_agent_engine( else: region = env_region click.echo(f'{region=} set by GOOGLE_CLOUD_LOCATION in {env_file}') + if api_key: + if 'GOOGLE_API_KEY' in env_vars: + click.secho( + 'Ignoring GOOGLE_API_KEY in .env as `--api_key` was' + ' explicitly passed and takes precedence', + fg='yellow', + ) + else: + env_vars['GOOGLE_GENAI_USE_VERTEXAI'] = '1' + env_vars['GOOGLE_API_KEY'] = api_key + elif not project: + if 'GOOGLE_API_KEY' in env_vars: + api_key = env_vars['GOOGLE_API_KEY'] + click.echo(f'api_key set by GOOGLE_API_KEY in {env_file}') if env_vars: if 'env_vars' in agent_config: click.echo( @@ -527,11 +843,20 @@ def to_agent_engine( # Set env_vars in agent_config to None if it is not set. agent_config['env_vars'] = agent_config.get('env_vars', env_vars) - vertexai.init( - project=project, - location=region, - staging_bucket=staging_bucket, - ) + import vertexai + + if project and region: + click.echo('Initializing Vertex AI...') + client = vertexai.Client(project=project, location=region) + elif api_key: + click.echo('Initializing Vertex AI in Express Mode with API key...') + client = vertexai.Client(api_key=api_key) + else: + click.echo( + 'No project/region or api_key provided. ' + 'Please specify either project/region or api_key.' + ) + return click.echo('Vertex AI initialized.') is_config_agent = False @@ -541,6 +866,16 @@ def to_agent_engine( is_config_agent = True adk_app_file = os.path.join(temp_folder, f'{adk_app}.py') + if adk_app_object == 'root_agent': + adk_app_type = 'agent' + elif adk_app_object == 'app': + adk_app_type = 'app' + else: + click.echo( + f'Invalid adk_app_object: {adk_app_object}. Please use "root_agent"' + ' or "app".' + ) + return with open(adk_app_file, 'w', encoding='utf-8') as f: f.write( _AGENT_ENGINE_APP_TEMPLATE.format( @@ -549,55 +884,36 @@ def to_agent_engine( is_config_agent=is_config_agent, temp_folder=temp_folder, agent_folder=agent_folder, + adk_app_object=adk_app_object, + adk_app_type=adk_app_type, + express_mode=api_key is not None, ) ) 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( + 'Agent Engine deployments have switched to source-based deployment, ' + 'so it is no longer necessary to absolutize imports.' + ) click.echo('Deploying to agent engine...') - agent_config['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_framework='google-adk', - ) + agent_config['entrypoint_module'] = f'{temp_folder}.{adk_app}' + agent_config['entrypoint_object'] = 'adk_app' + agent_config['source_packages'] = [temp_folder] + agent_config['class_methods'] = _AGENT_ENGINE_CLASS_METHODS + agent_config['agent_framework'] = 'google-adk' if not agent_engine_id: - agent_engines.create(**agent_config) + agent_engine = client.agent_engines.create(config=agent_config) + click.secho( + f'✅ Created agent engine: {agent_engine.api_resource.name}', + fg='green', + ) else: - resource_name = f'projects/{project}/locations/{region}/reasoningEngines/{agent_engine_id}' - agent_engines.update(resource_name=resource_name, **agent_config) + if project and region and not agent_engine_id.startswith('projects/'): + agent_engine_id = f'projects/{project}/locations/{region}/agentEngines/{agent_engine_id}' + client.agent_engines.update(name=agent_engine_id, config=agent_config) + click.secho(f'✅ Updated agent engine: {agent_engine_id}', fg='green') 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 06f567a0..5a71384b 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/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, diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index b2a31f70..696344eb 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/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,