13
13
# limitations under the License.
14
14
from __future__ import annotations
15
15
16
+ import json
16
17
import os
17
18
import shutil
18
19
import subprocess
@@ -348,6 +349,7 @@ def to_agent_engine(
348
349
description : Optional [str ] = None ,
349
350
requirements_file : Optional [str ] = None ,
350
351
env_file : Optional [str ] = None ,
352
+ agent_engine_config_file : Optional [str ] = None ,
351
353
):
352
354
"""Deploys an agent to Vertex AI Agent Engine.
353
355
@@ -397,6 +399,9 @@ def to_agent_engine(
397
399
variables. If not specified, the `.env` file in the `agent_folder` will be
398
400
used. The values of `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION`
399
401
will be overridden by `project` and `region` if they are specified.
402
+ agent_engine_config_file (str): The filepath to the agent engine config file
403
+ to use. If not specified, the `.agent_engine_config.json` file in the
404
+ `agent_folder` will be used.
400
405
"""
401
406
app_name = os .path .basename (agent_folder )
402
407
agent_src_path = os .path .join (temp_folder , app_name )
@@ -427,6 +432,34 @@ def to_agent_engine(
427
432
project = _resolve_project (project )
428
433
429
434
click .echo ('Resolving files and dependencies...' )
435
+ agent_config = {}
436
+ if not agent_engine_config_file :
437
+ # Attempt to read the agent engine config from .agent_engine_config.json in the dir (if any).
438
+ agent_engine_config_file = os .path .join (
439
+ agent_folder , '.agent_engine_config.json'
440
+ )
441
+ if os .path .exists (agent_engine_config_file ):
442
+ click .echo (f'Reading agent engine config from { agent_engine_config_file } ' )
443
+ with open (agent_engine_config_file , 'r' ) as f :
444
+ agent_config = json .load (f )
445
+ if display_name :
446
+ if 'display_name' in agent_config :
447
+ click .echo (
448
+ 'Overriding display_name in agent engine config with'
449
+ f' { display_name } '
450
+ )
451
+ agent_config ['display_name' ] = display_name
452
+ if description :
453
+ if 'description' in agent_config :
454
+ click .echo (
455
+ f'Overriding description in agent engine config with { description } '
456
+ )
457
+ agent_config ['description' ] = description
458
+ if agent_config .get ('extra_packages' ):
459
+ agent_config ['extra_packages' ].append (temp_folder )
460
+ else :
461
+ agent_config ['extra_packages' ] = [temp_folder ]
462
+
430
463
if not requirements_file :
431
464
# Attempt to read requirements from requirements.txt in the dir (if any).
432
465
requirements_txt_path = os .path .join (agent_src_path , 'requirements.txt' )
@@ -435,7 +468,18 @@ def to_agent_engine(
435
468
with open (requirements_txt_path , 'w' , encoding = 'utf-8' ) as f :
436
469
f .write ('google-cloud-aiplatform[adk,agent_engines]' )
437
470
click .echo (f'Created { requirements_txt_path } ' )
438
- requirements_file = requirements_txt_path
471
+ agent_config ['requirements' ] = agent_config .get (
472
+ 'requirements' ,
473
+ requirements_txt_path ,
474
+ )
475
+ else :
476
+ if 'requirements' in agent_config :
477
+ click .echo (
478
+ 'Overriding requirements in agent engine config with '
479
+ f'{ requirements_file } '
480
+ )
481
+ agent_config ['requirements' ] = requirements_file
482
+
439
483
env_vars = None
440
484
if not env_file :
441
485
# Attempt to read the env variables from .env in the dir (if any).
@@ -469,6 +513,14 @@ def to_agent_engine(
469
513
else :
470
514
region = env_region
471
515
click .echo (f'{ region = } set by GOOGLE_CLOUD_LOCATION in { env_file } ' )
516
+ if env_vars :
517
+ if 'env_vars' in agent_config :
518
+ click .echo (
519
+ f'Overriding env_vars in agent engine config with { env_vars } '
520
+ )
521
+ agent_config ['env_vars' ] = env_vars
522
+ # Set env_vars in agent_config to None if it is not set.
523
+ agent_config ['env_vars' ] = agent_config .get ('env_vars' , env_vars )
472
524
473
525
vertexai .init (
474
526
project = project ,
@@ -480,7 +532,7 @@ def to_agent_engine(
480
532
is_config_agent = False
481
533
config_root_agent_file = os .path .join (agent_src_path , 'root_agent.yaml' )
482
534
if os .path .exists (config_root_agent_file ):
483
- click .echo ('Config agent detected. ' )
535
+ click .echo (f 'Config agent detected: { config_root_agent_file } ' )
484
536
is_config_agent = True
485
537
486
538
adk_app_file = os .path .join (temp_folder , f'{ adk_app } .py' )
@@ -513,7 +565,7 @@ def to_agent_engine(
513
565
click .echo (f'The following exception was raised: { e } ' )
514
566
515
567
click .echo ('Deploying to agent engine...' )
516
- agent_engine = agent_engines .ModuleAgent (
568
+ agent_config [ ' agent_engine' ] = agent_engines .ModuleAgent (
517
569
module_name = adk_app ,
518
570
agent_name = 'adk_app' ,
519
571
register_operations = {
@@ -535,14 +587,6 @@ def to_agent_engine(
535
587
sys_paths = [temp_folder [1 :]],
536
588
agent_framework = 'google-adk' ,
537
589
)
538
- agent_config = dict (
539
- agent_engine = agent_engine ,
540
- requirements = requirements_file ,
541
- display_name = display_name ,
542
- description = description ,
543
- env_vars = env_vars ,
544
- extra_packages = [temp_folder ],
545
- )
546
590
547
591
if not agent_engine_id :
548
592
agent_engines .create (** agent_config )
0 commit comments