feat: add support for BYO-dockerfile in AE deployment

PiperOrigin-RevId: 865065204

Author: yeesian

vertexai/_genai/agent_engines.py

@@ -956,6 +956,7 @@ def create(
             agent_framework=config.agent_framework,
             python_version=config.python_version,
             build_options=config.build_options,
+            image_spec=config.image_spec,
         )
         operation = self._create(config=api_config)
         reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
@@ -1012,6 +1013,9 @@ def _set_source_code_spec(
         requirements_file: Optional[str] = None,
         sys_version: str,
         build_options: Optional[dict[str, list[str]]] = None,
+        image_spec: Optional[
+            types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
+        ] = None,
     ) -> None:
         """Sets source_code_spec for agent engine inside the `spec`."""
         source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict()
@@ -1035,7 +1039,25 @@ def _set_source_code_spec(
             raise ValueError(
                 "Please specify one of `source_packages` or `developer_connect_source`."
             )
-
+        if class_methods is None:
+            raise ValueError(
+                "`class_methods` must be specified if `source_packages` or `developer_connect_source` is specified."
+            )
+        update_masks.append("spec.class_methods")
+        class_methods_spec_list = (
+            _agent_engines_utils._class_methods_to_class_methods_spec(
+                class_methods=class_methods
+            )
+        )
+        spec["class_methods"] = [
+            _agent_engines_utils._to_dict(class_method_spec)
+            for class_method_spec in class_methods_spec_list
+        ]
+        if image_spec is not None:
+            update_masks.append("spec.source_code_spec.image_spec")
+            source_code_spec["image_spec"] = image_spec
+            spec["source_code_spec"] = source_code_spec
+            return
         update_masks.append("spec.source_code_spec.python_spec.version")
         python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = {
             "version": sys_version,
@@ -1058,21 +1080,6 @@ def _set_source_code_spec(
         source_code_spec["python_spec"] = python_spec
         spec["source_code_spec"] = source_code_spec
 
-        if class_methods is None:
-            raise ValueError(
-                "`class_methods` must be specified if `source_packages` or `developer_connect_source` is specified."
-            )
-        update_masks.append("spec.class_methods")
-        class_methods_spec_list = (
-            _agent_engines_utils._class_methods_to_class_methods_spec(
-                class_methods=class_methods
-            )
-        )
-        spec["class_methods"] = [
-            _agent_engines_utils._to_dict(class_method_spec)
-            for class_method_spec in class_methods_spec_list
-        ]
-
     def _set_package_spec(
         self,
         *,
@@ -1200,6 +1207,9 @@ def _create_config(
         agent_framework: Optional[str] = None,
         python_version: Optional[str] = None,
         build_options: Optional[dict[str, list[str]]] = None,
+        image_spec: Optional[
+            types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
+        ] = None,
     ) -> types.UpdateAgentEngineConfigDict:
         import sys
 
@@ -1285,6 +1295,7 @@ def _create_config(
                 requirements_file=requirements_file,
                 sys_version=sys_version,
                 build_options=build_options,
+                image_spec=image_spec,
             )
 
         if agent_engine_spec is not None:

vertexai/_genai/types/__init__.py

@@ -755,6 +755,9 @@
 from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict
 from .common import ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict
 from .common import ReasoningEngineSpecSourceCodeSpecDict
+from .common import ReasoningEngineSpecSourceCodeSpecImageSpec
+from .common import ReasoningEngineSpecSourceCodeSpecImageSpecDict
+from .common import ReasoningEngineSpecSourceCodeSpecImageSpecOrDict
 from .common import ReasoningEngineSpecSourceCodeSpecInlineSource
 from .common import ReasoningEngineSpecSourceCodeSpecInlineSourceDict
 from .common import ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict
@@ -1416,6 +1419,9 @@
     "ReasoningEngineSpecSourceCodeSpecPythonSpec",
     "ReasoningEngineSpecSourceCodeSpecPythonSpecDict",
     "ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict",
+    "ReasoningEngineSpecSourceCodeSpecImageSpec",
+    "ReasoningEngineSpecSourceCodeSpecImageSpecDict",
+    "ReasoningEngineSpecSourceCodeSpecImageSpecOrDict",
     "ReasoningEngineSpecSourceCodeSpec",
     "ReasoningEngineSpecSourceCodeSpecDict",
     "ReasoningEngineSpecSourceCodeSpecOrDict",

vertexai/_genai/types/common.py

@@ -5118,6 +5118,34 @@ class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False):
 ]
 
 
+class ReasoningEngineSpecSourceCodeSpecImageSpec(_common.BaseModel):
+    """The image spec for building an image (within a single build step).
+
+    It is based on the config file (i.e. Dockerfile) in the source directory.
+    """
+
+    build_args: Optional[dict[str, str]] = Field(
+        default=None,
+        description="""Optional. Build arguments to be used. They will be passed through --build-arg flags.""",
+    )
+
+
+class ReasoningEngineSpecSourceCodeSpecImageSpecDict(TypedDict, total=False):
+    """The image spec for building an image (within a single build step).
+
+    It is based on the config file (i.e. Dockerfile) in the source directory.
+    """
+
+    build_args: Optional[dict[str, str]]
+    """Optional. Build arguments to be used. They will be passed through --build-arg flags."""
+
+
+ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[
+    ReasoningEngineSpecSourceCodeSpecImageSpec,
+    ReasoningEngineSpecSourceCodeSpecImageSpecDict,
+]
+
+
 class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel):
     """Specification for deploying from source code."""
 
@@ -5133,6 +5161,10 @@ class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel):
     python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpec] = Field(
         default=None, description="""Configuration for a Python application."""
     )
+    image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field(
+        default=None,
+        description="""Optional. Configuration for building an image with custom config file.""",
+    )
 
 
 class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False):
@@ -5149,6 +5181,9 @@ class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False):
     python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict]
     """Configuration for a Python application."""
 
+    image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict]
+    """Optional. Configuration for building an image with custom config file."""
+
 
 ReasoningEngineSpecSourceCodeSpecOrDict = Union[
     ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict
@@ -14004,6 +14039,9 @@ class AgentEngineConfig(_common.BaseModel):
           subdirectory and the path must be added to `extra_packages`.
       """,
     )
+    image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field(
+        default=None, description="""The image spec for the Agent Engine."""
+    )
 
 
 class AgentEngineConfigDict(TypedDict, total=False):
@@ -14167,6 +14205,9 @@ class AgentEngineConfigDict(TypedDict, total=False):
           subdirectory and the path must be added to `extra_packages`.
       """
 
+    image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict]
+    """The image spec for the Agent Engine."""
+
 
 AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]