feat: GenAI SDK client - Support build options in Agent Engine source-based Deployment.

PiperOrigin-RevId: 833525372

Author: shawn-yang-google

tests/unit/vertexai/genai/replays/conftest.py

@@ -123,6 +123,13 @@ def replays_prefix():
     return "test"
 
 
+@pytest.fixture
+def mock_agent_engine_create_path_exists():
+    """Mocks os.path.exists to return True."""
+    with mock.patch("os.path.exists", return_value=True) as mock_exists:
+        yield mock_exists
+
+
 @pytest.fixture
 def mock_agent_engine_create_base64_encoded_tarball():
     """Mocks the _create_base64_encoded_tarball function."""

tests/unit/vertexai/genai/replays/test_create_agent_engine.py

@@ -115,9 +115,13 @@ def test_create_with_context_spec(client):
 def test_create_with_source_packages(
     client,
     mock_agent_engine_create_base64_encoded_tarball,
+    mock_agent_engine_create_path_exists,
 ):
     """Tests creating an agent engine with source packages."""
-    with mock_agent_engine_create_base64_encoded_tarball:
+    with (
+        mock_agent_engine_create_base64_encoded_tarball,
+        mock_agent_engine_create_path_exists,
+    ):
         agent_engine = client.agent_engines.create(
             config={
                 "display_name": "test-agent-engine-source-packages",

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1062,10 +1062,42 @@ def test_create_agent_engine_config_with_source_packages(
                 == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
             )
 
+    @mock.patch.object(
+        _agent_engines_utils,
+        "_create_base64_encoded_tarball",
+        return_value="test_tarball",
+    )
+    @mock.patch.object(_agent_engines_utils, "_validate_packages_or_raise")
+    def test_create_agent_engine_config_with_source_packages_and_build_options(
+        self, mock_validate_packages, mock_create_base64_encoded_tarball
+    ):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_file_path = os.path.join(tmpdir, "main.py")
+            with open(test_file_path, "w") as f:
+                f.write("app = None")
+            build_options = {
+                "installation_scripts": ["installation_scripts/install.sh"]
+            }
+            source_packages = [test_file_path, "installation_scripts/install.sh"]
+            mock_validate_packages.return_value = source_packages
+
+            self.client.agent_engines._create_config(
+                mode="create",
+                source_packages=source_packages,
+                entrypoint_module="main",
+                entrypoint_object="app",
+                build_options=build_options,
+                class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+            )
+            mock_validate_packages.assert_called_once_with(
+                packages=source_packages,
+                build_options=build_options,
+            )
+
     @mock.patch.object(_agent_engines_utils, "_prepare")
-    @mock.patch.object(_agent_engines_utils, "_validate_extra_packages_or_raise")
+    @mock.patch.object(_agent_engines_utils, "_validate_packages_or_raise")
     def test_create_agent_engine_config_with_build_options(
-        self, mock_validate_extra_packages, mock_prepare
+        self, mock_validate_packages, mock_prepare
     ):
         build_options = {"installation_scripts": ["install.sh"]}
         extra_packages = ["install.sh"]
@@ -1079,8 +1111,8 @@ def test_create_agent_engine_config_with_build_options(
             build_options=build_options,
         )
 
-        mock_validate_extra_packages.assert_called_once_with(
-            extra_packages=extra_packages,
+        mock_validate_packages.assert_called_once_with(
+            packages=extra_packages,
             build_options=build_options,
         )
 

vertexai/_genai/_agent_engines_utils.py

@@ -1238,42 +1238,40 @@ def _create_base64_encoded_tarball(
     return base64.b64encode(tarball_bytes).decode("utf-8")
 
 
-def _validate_extra_packages_or_raise(
+def _validate_packages_or_raise(
     *,
-    extra_packages: Sequence[str],
+    packages: Sequence[str],
     build_options: Optional[Dict[str, Sequence[str]]] = None,
 ) -> Sequence[str]:
-    """Tries to validates the extra packages."""
-    extra_packages = extra_packages or []
+    """Tries to validates the packages."""
+    packages = packages or []
     if build_options and _INSTALLATION_SUBDIR in build_options:
         _validate_installation_scripts_or_raise(
             script_paths=build_options[_INSTALLATION_SUBDIR],
-            extra_packages=extra_packages,
+            packages=packages,
         )
-    for extra_package in extra_packages:
-        if not os.path.exists(extra_package):
-            raise FileNotFoundError(
-                f"Extra package specified but not found: {extra_package=}"
-            )
-    return extra_packages
+    for package in packages:
+        if not os.path.exists(package):
+            raise FileNotFoundError(f"Package specified but not found: {package=}")
+    return packages
 
 
 def _validate_installation_scripts_or_raise(
     *,
     script_paths: Sequence[str],
-    extra_packages: Sequence[str],
+    packages: Sequence[str],
 ) -> None:
     """Validates the installation scripts' path explicitly provided by the user.
 
     Args:
         script_paths (Sequence[str]):
             Required. The paths to the installation scripts.
-        extra_packages (Sequence[str]):
-            Required. The extra packages to be updated.
+        packages (Sequence[str]):
+            Required. The user-provided packages.
 
     Raises:
         ValueError: If a user-defined script is not under the expected
-            subdirectory, or not in `extra_packages`, or if an extra package is
+            subdirectory, or not in `packages`, or if a package is
             in the installation scripts subdirectory, but is not specified as an
             installation script.
     """
@@ -1283,37 +1281,35 @@ def _validate_installation_scripts_or_raise(
                 f"User-defined installation script '{script_path}' is not in "
                 f"the expected '{_INSTALLATION_SUBDIR}' subdirectory. "
                 f"Ensure it is placed in '{_INSTALLATION_SUBDIR}' within your "
-                f"`extra_packages`."
+                f"'extra_packages' or 'source_packages'."
             )
             raise ValueError(
                 f"Required installation script '{script_path}' "
                 f"is not under '{_INSTALLATION_SUBDIR}'"
             )
 
-        if script_path not in extra_packages:
+        if script_path not in packages:
             logger.warning(
                 f"User-defined installation script '{script_path}' is not in "
-                f"extra_packages. Ensure it is added to `extra_packages`."
+                f"'extra_packages' or 'source_packages'. Ensure it is added to "
+                f"'extra_packages' or 'source_packages'."
             )
             raise ValueError(
                 f"User-defined installation script '{script_path}' "
-                f"does not exist in `extra_packages`"
+                f"does not exist in 'extra_packages' or 'source_packages'."
             )
 
-    for extra_package in extra_packages:
-        if (
-            extra_package.startswith(_INSTALLATION_SUBDIR)
-            and extra_package not in script_paths
-        ):
+    for package in packages:
+        if package.startswith(_INSTALLATION_SUBDIR) and package not in script_paths:
             logger.warning(
-                f"Extra package '{extra_package}' is in the installation "
+                f"Package '{package}' is in the installation "
                 "scripts subdirectory, but is not specified as an installation "
                 "script in `build_options`. "
                 "Ensure it is added to installation_scripts for "
                 "automatic execution."
             )
             raise ValueError(
-                f"Extra package '{extra_package}' is in the installation "
+                f"Package '{package}' is in the installation "
                 "scripts subdirectory, but is not specified as an installation "
                 "script in `build_options`."
             )

vertexai/_genai/agent_engines.py

@@ -1067,8 +1067,8 @@ def _create_config(
                 agent=agent,
                 requirements=requirements,
             )
-            extra_packages = _agent_engines_utils._validate_extra_packages_or_raise(
-                extra_packages=extra_packages,
+            extra_packages = _agent_engines_utils._validate_packages_or_raise(
+                packages=extra_packages,
                 build_options=build_options,
             )
             # Prepares the Agent Engine for creation/update in Vertex AI. This
@@ -1135,6 +1135,10 @@ def _create_config(
             }
 
         if source_packages is not None:
+            source_packages = _agent_engines_utils._validate_packages_or_raise(
+                packages=source_packages,
+                build_options=build_options,
+            )
             update_masks.append("spec.source_code_spec.inline_source.source_archive")
             source_code_spec = {
                 "inline_source": {