From ef03befdd4c4a91c3b70d8ebed2d33d3a9c85a03 Mon Sep 17 00:00:00 2001 From: Shawn Yang Date: Tue, 11 Nov 2025 12:05:44 -0800 Subject: [PATCH] chore: Support specifying python_version in Agent Engine creation and update. PiperOrigin-RevId: 831015034 --- .../unit/vertexai/genai/test_agent_engines.py | 15 +++++-- vertexai/_genai/agent_engines.py | 14 ++++++- vertexai/_genai/types/common.py | 39 +++++++++++++++++++ 3 files changed, 64 insertions(+), 4 deletions(-) diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py index e74fcf2607..afd9acadb4 100644 --- a/tests/unit/vertexai/genai/test_agent_engines.py +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -458,6 +458,7 @@ def register_operations(self) -> Dict[str, List[str]]: ) _TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE _TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}" +_TEST_PYTHON_VERSION_OVERRIDE = "3.11" _TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK _TEST_AGENT_ENGINE_CLASS_METHOD_1 = { "description": "Runs the engine.", @@ -969,12 +970,13 @@ def test_create_agent_engine_config_full(self, mock_prepare): resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS, container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY, encryption_spec=_TEST_AGENT_ENGINE_ENCRYPTION_SPEC, + python_version=_TEST_PYTHON_VERSION_OVERRIDE, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION assert config["spec"]["agent_framework"] == "custom" assert config["spec"]["package_spec"] == { - "python_version": _TEST_PYTHON_VERSION, + "python_version": _TEST_PYTHON_VERSION_OVERRIDE, "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, "dependency_files_gcs_uri": _TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI, "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, @@ -1037,6 +1039,7 @@ def test_create_agent_engine_config_with_source_packages( class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, agent_framework=_TEST_AGENT_FRAMEWORK, identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, + python_version=_TEST_PYTHON_VERSION_OVERRIDE, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION @@ -1044,7 +1047,7 @@ def test_create_agent_engine_config_with_source_packages( assert config["spec"]["source_code_spec"] == { "inline_source": {"source_archive": "test_tarball"}, "python_spec": { - "version": _TEST_PYTHON_VERSION, + "version": _TEST_PYTHON_VERSION_OVERRIDE, "entrypoint_module": "main", "entrypoint_object": "app", "requirements_file": requirements_file_path, @@ -1073,12 +1076,13 @@ def test_update_agent_engine_config_full(self, mock_prepare): env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT, identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, + python_version=_TEST_PYTHON_VERSION_OVERRIDE, ) assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION assert config["spec"]["agent_framework"] == "custom" assert config["spec"]["package_spec"] == { - "python_version": _TEST_PYTHON_VERSION, + "python_version": _TEST_PYTHON_VERSION_OVERRIDE, "pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI, "dependency_files_gcs_uri": _TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI, "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, @@ -1593,6 +1597,7 @@ def test_create_agent_engine_with_env_vars_dict( entrypoint_object=None, requirements_file=None, agent_framework=None, + python_version=None, ) request_mock.assert_called_with( "post", @@ -1685,6 +1690,7 @@ def test_create_agent_engine_with_custom_service_account( entrypoint_object=None, requirements_file=None, agent_framework=None, + python_version=None, ) request_mock.assert_called_with( "post", @@ -1776,6 +1782,7 @@ def test_create_agent_engine_with_experimental_mode( entrypoint_object=None, requirements_file=None, agent_framework=None, + python_version=None, ) request_mock.assert_called_with( "post", @@ -1930,6 +1937,7 @@ def test_create_agent_engine_with_class_methods( entrypoint_object=None, requirements_file=None, agent_framework=None, + python_version=None, ) request_mock.assert_called_with( "post", @@ -2016,6 +2024,7 @@ def test_create_agent_engine_with_agent_framework( requirements_file=None, agent_framework=_TEST_AGENT_FRAMEWORK, identity_type=None, + python_version=None, ) request_mock.assert_called_with( "post", diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py index 5d57c6fd70..22d155e989 100644 --- a/vertexai/_genai/agent_engines.py +++ b/vertexai/_genai/agent_engines.py @@ -97,6 +97,9 @@ def _CreateAgentEngineConfig_to_vertex( if getv(from_object, ["agent_framework"]) is not None: setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + return to_object @@ -291,6 +294,9 @@ def _UpdateAgentEngineConfig_to_vertex( if getv(from_object, ["agent_framework"]) is not None: setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + if getv(from_object, ["update_mask"]) is not None: setv( parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) @@ -931,6 +937,7 @@ def create( entrypoint_object=config.entrypoint_object, requirements_file=config.requirements_file, agent_framework=config.agent_framework, + python_version=config.python_version, ) operation = self._create(config=api_config) # TODO: Use a more specific link. @@ -996,6 +1003,7 @@ def _create_config( entrypoint_object: Optional[str] = None, requirements_file: Optional[str] = None, agent_framework: Optional[str] = None, + python_version: Optional[str] = None, ) -> types.UpdateAgentEngineConfigDict: import sys @@ -1027,7 +1035,10 @@ def _create_config( if agent_framework == "google-adk": env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) - sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" + if python_version: + sys_version = python_version + else: + sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" agent_engine_spec = None if agent is not None: if source_packages is not None: @@ -1453,6 +1464,7 @@ def update( entrypoint_object=config.entrypoint_object, requirements_file=config.requirements_file, agent_framework=config.agent_framework, + python_version=config.python_version, ) operation = self._update(name=name, config=api_config) logger.info( diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 34619b654e..5a2e2dbbf9 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -5457,6 +5457,13 @@ class CreateAgentEngineConfig(_common.BaseModel): - If `source_packages` is specified, the agent framework will default to "custom".""", ) + python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] = Field( + default=None, + description="""The Python version to be used for the Agent Engine. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13". + """, + ) class CreateAgentEngineConfigDict(TypedDict, total=False): @@ -5567,6 +5574,12 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): - If `source_packages` is specified, the agent framework will default to "custom".""" + python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] + """The Python version to be used for the Agent Engine. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13". + """ + CreateAgentEngineConfigOrDict = Union[ CreateAgentEngineConfig, CreateAgentEngineConfigDict @@ -6183,6 +6196,13 @@ class UpdateAgentEngineConfig(_common.BaseModel): - If `source_packages` is specified, the agent framework will default to "custom".""", ) + python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] = Field( + default=None, + description="""The Python version to be used for the Agent Engine. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13". + """, + ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -6298,6 +6318,12 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): - If `source_packages` is specified, the agent framework will default to "custom".""" + python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] + """The Python version to be used for the Agent Engine. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13". + """ + update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -13037,6 +13063,13 @@ class AgentEngineConfig(_common.BaseModel): - If `source_packages` is specified, the agent framework will default to "custom".""", ) + python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] = Field( + default=None, + description="""The Python version to be used for the Agent Engine. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13". + """, + ) class AgentEngineConfigDict(TypedDict, total=False): @@ -13179,6 +13212,12 @@ class AgentEngineConfigDict(TypedDict, total=False): - If `source_packages` is specified, the agent framework will default to "custom".""" + python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] + """The Python version to be used for the Agent Engine. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13". + """ + AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]