diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 0b4d8ecb9a..51998f12cd 100644 --- a/vertexai/_genai/agent_engines.py +++ b/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: diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index fa43a483dc..9565534e14 100644 --- a/vertexai/_genai/types/__init__.py +++ b/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", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 2716964936..1d46d80a80 100644 --- a/vertexai/_genai/types/common.py +++ b/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]