Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions tests/unit/vertexai/genai/test_agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -1062,6 +1062,28 @@ def test_create_agent_engine_config_with_source_packages(
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
)

@mock.patch.object(_agent_engines_utils, "_prepare")
@mock.patch.object(_agent_engines_utils, "_validate_extra_packages_or_raise")
def test_create_agent_engine_config_with_build_options(
self, mock_validate_extra_packages, mock_prepare
):
build_options = {"installation_scripts": ["install.sh"]}
extra_packages = ["install.sh"]

self.client.agent_engines._create_config(
mode="create",
agent=self.test_agent,
staging_bucket=_TEST_STAGING_BUCKET,
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
extra_packages=extra_packages,
build_options=build_options,
)

mock_validate_extra_packages.assert_called_once_with(
extra_packages=extra_packages,
build_options=build_options,
)

@mock.patch.object(_agent_engines_utils, "_prepare")
def test_update_agent_engine_config_full(self, mock_prepare):
config = self.client.agent_engines._create_config(
Expand Down Expand Up @@ -1598,6 +1620,7 @@ def test_create_agent_engine_with_env_vars_dict(
requirements_file=None,
agent_framework=None,
python_version=None,
build_options=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -1691,6 +1714,7 @@ def test_create_agent_engine_with_custom_service_account(
requirements_file=None,
agent_framework=None,
python_version=None,
build_options=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -1783,6 +1807,7 @@ def test_create_agent_engine_with_experimental_mode(
requirements_file=None,
agent_framework=None,
python_version=None,
build_options=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -1938,6 +1963,7 @@ def test_create_agent_engine_with_class_methods(
requirements_file=None,
agent_framework=None,
python_version=None,
build_options=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2025,6 +2051,7 @@ def test_create_agent_engine_with_agent_framework(
agent_framework=_TEST_AGENT_FRAMEWORK,
identity_type=None,
python_version=None,
build_options=None,
)
request_mock.assert_called_with(
"post",
Expand Down
4 changes: 4 additions & 0 deletions vertexai/_genai/agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -938,6 +938,7 @@ def create(
requirements_file=config.requirements_file,
agent_framework=config.agent_framework,
python_version=config.python_version,
build_options=config.build_options,
)
operation = self._create(config=api_config)
# TODO: Use a more specific link.
Expand Down Expand Up @@ -1004,6 +1005,7 @@ def _create_config(
requirements_file: Optional[str] = None,
agent_framework: Optional[str] = None,
python_version: Optional[str] = None,
build_options: Optional[dict[str, list[str]]] = None,
) -> types.UpdateAgentEngineConfigDict:
import sys

Expand Down Expand Up @@ -1067,6 +1069,7 @@ def _create_config(
)
extra_packages = _agent_engines_utils._validate_extra_packages_or_raise(
extra_packages=extra_packages,
build_options=build_options,
)
# Prepares the Agent Engine for creation/update in Vertex AI. This
# involves packaging and uploading the artifacts for agent_engine,
Expand Down Expand Up @@ -1465,6 +1468,7 @@ def update(
requirements_file=config.requirements_file,
agent_framework=config.agent_framework,
python_version=config.python_version,
build_options=config.build_options,
)
operation = self._update(name=name, config=api_config)
logger.info(
Expand Down
63 changes: 63 additions & 0 deletions vertexai/_genai/types/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -5471,6 +5471,17 @@ class CreateAgentEngineConfig(_common.BaseModel):
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
""",
)
build_options: Optional[dict[str, list[str]]] = Field(
default=None,
description="""The build options for the Agent Engine.
The following keys are supported:
- installation_scripts:
Optional. The paths to the installation scripts to be
executed in the Docker image.
The scripts must be located in the `installation_scripts`
subdirectory and the path must be added to `extra_packages`.
""",
)


class CreateAgentEngineConfigDict(TypedDict, total=False):
Expand Down Expand Up @@ -5587,6 +5598,16 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
"""

build_options: Optional[dict[str, list[str]]]
"""The build options for the Agent Engine.
The following keys are supported:
- installation_scripts:
Optional. The paths to the installation scripts to be
executed in the Docker image.
The scripts must be located in the `installation_scripts`
subdirectory and the path must be added to `extra_packages`.
"""


CreateAgentEngineConfigOrDict = Union[
CreateAgentEngineConfig, CreateAgentEngineConfigDict
Expand Down Expand Up @@ -6210,6 +6231,17 @@ class UpdateAgentEngineConfig(_common.BaseModel):
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
""",
)
build_options: Optional[dict[str, list[str]]] = Field(
default=None,
description="""The build options for the Agent Engine.
The following keys are supported:
- installation_scripts:
Optional. The paths to the installation scripts to be
executed in the Docker image.
The scripts must be located in the `installation_scripts`
subdirectory and the path must be added to `extra_packages`.
""",
)
update_mask: Optional[str] = Field(
default=None,
description="""The update mask to apply. For the `FieldMask` definition, see
Expand Down Expand Up @@ -6331,6 +6363,16 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
"""

build_options: Optional[dict[str, list[str]]]
"""The build options for the Agent Engine.
The following keys are supported:
- installation_scripts:
Optional. The paths to the installation scripts to be
executed in the Docker image.
The scripts must be located in the `installation_scripts`
subdirectory and the path must be added to `extra_packages`.
"""

update_mask: Optional[str]
"""The update mask to apply. For the `FieldMask` definition, see
https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
Expand Down Expand Up @@ -13251,6 +13293,17 @@ class AgentEngineConfig(_common.BaseModel):
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
""",
)
build_options: Optional[dict[str, list[str]]] = Field(
default=None,
description="""The build options for the Agent Engine.
The following keys are supported:
- installation_scripts:
Optional. The paths to the installation scripts to be
executed in the Docker image.
The scripts must be located in the `installation_scripts`
subdirectory and the path must be added to `extra_packages`.
""",
)


class AgentEngineConfigDict(TypedDict, total=False):
Expand Down Expand Up @@ -13399,6 +13452,16 @@ class AgentEngineConfigDict(TypedDict, total=False):
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
"""

build_options: Optional[dict[str, list[str]]]
"""The build options for the Agent Engine.
The following keys are supported:
- installation_scripts:
Optional. The paths to the installation scripts to be
executed in the Docker image.
The scripts must be located in the `installation_scripts`
subdirectory and the path must be added to `extra_packages`.
"""


AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]

Expand Down
Loading