chore: Support specifying python_version in Agent Engine creation and update.

PiperOrigin-RevId: 831015034

Author: shawn-yang-google

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,14 +1039,15 @@ 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
             assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
             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",

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(

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]