Skip to content

Commit 3dc3f6d

Browse files
shawn-yang-googlecopybara-github
shawn-yang-google
authored and
copybara-github
committed
feat: **Allow installation scripts in AgentEngine.**
PiperOrigin-RevId: 770834726
1 parent c5bb99b commit 3dc3f6d

File tree

4 files changed

+303
-9
lines changed

4 files changed

+303
-9
lines changed

‎tests/unit/vertex_langchain/test_agent_engines.py

Lines changed: 172 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,12 @@ def register_operations(self) -> Dict[str, List[str]]:
585585
"pydantic": ["pydantic"],
586586
}
587587

588+
_TEST_BUILD_OPTIONS_INSTALLATION = _agent_engines._BUILD_OPTIONS_INSTALLATION
589+
_TEST_INSTALLATION_SUBDIR = _utils._INSTALLATION_SUBDIR
590+
_TEST_INSTALLATION_SCRIPT_PATH = (
591+
f"{_TEST_INSTALLATION_SUBDIR}/install_package.sh"
592+
)
593+
588594

589595
def _create_empty_fake_package(package_name: str) -> str:
590596
"""Creates a temporary directory structure representing an empty fake Python package.
@@ -1147,6 +1153,49 @@ def test_create_agent_engine_with_env_vars_list(
11471153
retry=_TEST_RETRY,
11481154
)
11491155

1156+
def test_create_agent_engine_with_build_options(
1157+
self,
1158+
create_agent_engine_mock,
1159+
cloud_storage_create_bucket_mock,
1160+
tarfile_open_mock,
1161+
cloudpickle_dump_mock,
1162+
cloudpickle_load_mock,
1163+
importlib_metadata_version_mock,
1164+
get_agent_engine_mock,
1165+
get_gca_resource_mock,
1166+
):
1167+
1168+
with mock.patch("os.path.exists", return_value=True):
1169+
agent_engines.create(
1170+
self.test_agent,
1171+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1172+
extra_packages=[
1173+
_TEST_INSTALLATION_SCRIPT_PATH,
1174+
],
1175+
build_options={
1176+
_TEST_BUILD_OPTIONS_INSTALLATION: [
1177+
_TEST_INSTALLATION_SCRIPT_PATH
1178+
]
1179+
},
1180+
)
1181+
test_spec = types.ReasoningEngineSpec(
1182+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
1183+
agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK,
1184+
)
1185+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
1186+
create_agent_engine_mock.assert_called_with(
1187+
parent=_TEST_PARENT,
1188+
reasoning_engine=types.ReasoningEngine(
1189+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1190+
spec=test_spec,
1191+
),
1192+
)
1193+
1194+
get_agent_engine_mock.assert_called_with(
1195+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1196+
retry=_TEST_RETRY,
1197+
)
1198+
11501199
# pytest does not allow absl.testing.parameterized.named_parameters.
11511200
@pytest.mark.parametrize(
11521201
"test_case_name, test_engine_instance, expected_framework",
@@ -2795,8 +2844,8 @@ def test_update_agent_engine_with_no_updates(
27952844
ValueError,
27962845
match=(
27972846
"At least one of `agent_engine`, `requirements`, "
2798-
"`extra_packages`, `display_name`, `description`, or `env_vars` "
2799-
"must be specified."
2847+
"`extra_packages`, `display_name`, `description`, "
2848+
"`env_vars`, or `build_options` must be specified."
28002849
),
28012850
):
28022851
test_agent_engine = _generate_agent_engine_to_update()
@@ -3227,3 +3276,124 @@ def test_scan_with_explicit_ignore_modules(self):
32273276
"cloudpickle": "3.0.0",
32283277
"pydantic": "1.11.1",
32293278
}
3279+
3280+
3281+
class TestValidateInstallationScripts:
3282+
@parameterized.named_parameters(
3283+
dict(
3284+
testcase_name="valid_script_in_subdir_and_extra_packages",
3285+
script_paths=[f"{_utils._INSTALLATION_SUBDIR}/script.sh"],
3286+
extra_packages=[f"{_utils._INSTALLATION_SUBDIR}/script.sh"],
3287+
raises_error=False,
3288+
),
3289+
dict(
3290+
testcase_name="script_not_in_subdir",
3291+
script_paths=["script.sh"],
3292+
extra_packages=["script.sh"],
3293+
raises_error=True,
3294+
error_message=(
3295+
f"Required installation script 'script.sh' must be under"
3296+
f" '{_utils._INSTALLATION_SUBDIR}'"
3297+
),
3298+
),
3299+
dict(
3300+
testcase_name="script_not_in_extra_packages",
3301+
script_paths=[f"{_utils._INSTALLATION_SUBDIR}/script.sh"],
3302+
extra_packages=[],
3303+
raises_error=True,
3304+
error_message=(
3305+
"Required installation script "
3306+
f"'{_utils._INSTALLATION_SUBDIR}/script.sh'"
3307+
" must be in extra_packages"
3308+
),
3309+
),
3310+
dict(
3311+
testcase_name="extra_package_in_subdir_but_not_script",
3312+
script_paths=[],
3313+
extra_packages=[f"{_utils._INSTALLATION_SUBDIR}/script.sh"],
3314+
raises_error=True,
3315+
error_message=(
3316+
f"Extra package '{_utils._INSTALLATION_SUBDIR}/script.sh' "
3317+
"is in the installation scripts subdirectory, but is not "
3318+
"specified as an installation script."
3319+
),
3320+
),
3321+
dict(
3322+
testcase_name="multiple_valid_scripts",
3323+
script_paths=[
3324+
f"{_utils._INSTALLATION_SUBDIR}/script1.sh",
3325+
f"{_utils._INSTALLATION_SUBDIR}/script2.sh",
3326+
],
3327+
extra_packages=[
3328+
f"{_utils._INSTALLATION_SUBDIR}/script1.sh",
3329+
f"{_utils._INSTALLATION_SUBDIR}/script2.sh",
3330+
],
3331+
raises_error=False,
3332+
),
3333+
dict(
3334+
testcase_name="one_valid_one_invalid_script_not_in_subdir",
3335+
script_paths=[
3336+
f"{_utils._INSTALLATION_SUBDIR}/script1.sh",
3337+
"script2.sh",
3338+
],
3339+
extra_packages=[
3340+
f"{_utils._INSTALLATION_SUBDIR}/script1.sh",
3341+
"script2.sh",
3342+
],
3343+
raises_error=True,
3344+
error_message=(
3345+
f"Required installation script 'script2.sh' must be under"
3346+
f" '{_utils._INSTALLATION_SUBDIR}'"
3347+
),
3348+
),
3349+
dict(
3350+
testcase_name="one_valid_one_invalid_script_not_in_extra_packages",
3351+
script_paths=[
3352+
f"{_utils._INSTALLATION_SUBDIR}/script1.sh",
3353+
f"{_utils._INSTALLATION_SUBDIR}/script2.sh",
3354+
],
3355+
extra_packages=[f"{_utils._INSTALLATION_SUBDIR}/script1.sh"],
3356+
raises_error=True,
3357+
error_message=(
3358+
f"Required installation script"
3359+
f" '{_utils._INSTALLATION_SUBDIR}/script2.sh' must be in"
3360+
f" extra_packages"
3361+
),
3362+
),
3363+
dict(
3364+
testcase_name="one_valid_one_invalid_extra_package_in_subdir",
3365+
script_paths=[f"{_utils._INSTALLATION_SUBDIR}/script1.sh"],
3366+
extra_packages=[
3367+
f"{_utils._INSTALLATION_SUBDIR}/script1.sh",
3368+
f"{_utils._INSTALLATION_SUBDIR}/script2.sh",
3369+
],
3370+
raises_error=True,
3371+
error_message=(
3372+
f"Extra package '{_utils._INSTALLATION_SUBDIR}/script2.sh' "
3373+
"is in the installation scripts subdirectory, but is not "
3374+
"specified as an installation script."
3375+
),
3376+
),
3377+
dict(
3378+
testcase_name="empty_script_paths_and_extra_packages",
3379+
script_paths=[],
3380+
extra_packages=[],
3381+
raises_error=False,
3382+
),
3383+
)
3384+
def test_validate_installation_scripts(
3385+
self,
3386+
script_paths,
3387+
extra_packages,
3388+
raises_error,
3389+
error_message=None,
3390+
):
3391+
if raises_error:
3392+
with pytest.raises(ValueError, match=error_message):
3393+
_utils.validate_installation_scripts_or_raise(
3394+
script_paths=script_paths, extra_packages=extra_packages
3395+
)
3396+
else:
3397+
_utils.validate_installation_scripts_or_raise(
3398+
script_paths=script_paths, extra_packages=extra_packages
3399+
)

‎vertexai/agent_engines/__init__.py

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ def create(
7070
env_vars: Optional[
7171
Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
7272
] = None,
73+
build_options: Optional[Dict[str, Sequence[str]]] = None,
7374
) -> AgentEngine:
7475
"""Creates a new Agent Engine.
7576
@@ -86,6 +87,9 @@ def create(
8687
|-- user_code/
8788
| |-- utils.py
8889
| |-- ...
90+
|-- installation_scripts/
91+
| |-- install_package.sh
92+
| |-- ...
8993
|-- ...
9094
9195
To build an Agent Engine with the above files, run:
@@ -105,6 +109,12 @@ def create(
105109
"./user_src_dir/user_code", # a directory
106110
...
107111
],
112+
build_options={
113+
"installation": [
114+
"./user_src_dir/installation_scripts/install_package.sh",
115+
...
116+
],
117+
},
108118
)
109119
110120
Args:
@@ -131,6 +141,9 @@ def create(
131141
a valid key to `os.environ`. If it is a dictionary, the keys are
132142
the environment variable names, and the values are the
133143
corresponding values.
144+
build_options (Dict[str, Sequence[str]]):
145+
Optional. The build options for the Agent Engine. This includes
146+
options such as installation scripts.
134147
135148
Returns:
136149
AgentEngine: The Agent Engine that was created.
@@ -153,6 +166,7 @@ def create(
153166
gcs_dir_name=gcs_dir_name,
154167
extra_packages=extra_packages,
155168
env_vars=env_vars,
169+
build_options=build_options,
156170
)
157171

158172

@@ -237,6 +251,7 @@ def update(
237251
env_vars: Optional[
238252
Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
239253
] = None,
254+
build_options: Optional[Dict[str, Sequence[str]]] = None,
240255
) -> "AgentEngine":
241256
"""Updates an existing Agent Engine.
242257
@@ -280,6 +295,9 @@ def update(
280295
a valid key to `os.environ`. If it is a dictionary, the keys are
281296
the environment variable names, and the values are the
282297
corresponding values.
298+
build_options (Dict[str, Sequence[str]]):
299+
Optional. The build options for the Agent Engine. This includes
300+
options such as installation scripts.
283301
284302
Returns:
285303
AgentEngine: The Agent Engine that was updated.
@@ -290,8 +308,8 @@ def update(
290308
FileNotFoundError: If `extra_packages` includes a file or directory
291309
that does not exist.
292310
ValueError: if none of `display_name`, `description`,
293-
`requirements`, `extra_packages`, or `agent_engine` were
294-
specified.
311+
`requirements`, `extra_packages`, `agent_engine`, or `build_options`
312+
were specified.
295313
IOError: If requirements is a string that corresponds to a
296314
nonexistent file.
297315
"""
@@ -304,6 +322,7 @@ def update(
304322
gcs_dir_name=gcs_dir_name,
305323
extra_packages=extra_packages,
306324
env_vars=env_vars,
325+
build_options=build_options,
307326
)
308327

309328

0 commit comments

Comments
 (0)