feat: Add the identity type option for the agent engine and add effective identity to the resource

PiperOrigin-RevId: 819733517

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_agent_engines.py

@@ -533,6 +533,9 @@ def register_operations(self) -> Dict[str, List[str]]:
 }
 _TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY = 4
 _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT = "test-custom-service-account"
+_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT = (
+    _genai_types.IdentityType.SERVICE_ACCOUNT
+)
 _TEST_AGENT_ENGINE_ENCRYPTION_SPEC = {"kms_key_name": "test-kms-key"}
 _TEST_AGENT_ENGINE_SPEC = _genai_types.ReasoningEngineSpecDict(
     agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK,
@@ -559,6 +562,7 @@ def register_operations(self) -> Dict[str, List[str]]:
         requirements_gcs_uri=_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
     ),
     service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
+    identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
 )
 _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}]
 _TEST_AGENT_ENGINE_OPERATION_SCHEMAS = []
@@ -908,6 +912,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
             extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
             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,
             psc_interface_config=_TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG,
             min_instances=_TEST_AGENT_ENGINE_MIN_INSTANCES,
             max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES,
@@ -950,6 +955,10 @@ def test_create_agent_engine_config_full(self, mock_prepare):
             config["spec"]["service_account"]
             == _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT
         )
+        assert (
+            config["spec"]["identity_type"]
+            == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+        )
 
     @mock.patch.object(
         _agent_engines_utils,
@@ -977,6 +986,7 @@ def test_create_agent_engine_config_with_source_packages(
                 requirements_file=requirements_file_path,
                 class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
                 agent_framework=_TEST_AGENT_FRAMEWORK,
+                identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
             )
             assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
             assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
@@ -994,6 +1004,10 @@ def test_create_agent_engine_config_with_source_packages(
             mock_create_base64_encoded_tarball.assert_called_once_with(
                 source_packages=[test_file_path]
             )
+            assert (
+                config["spec"]["identity_type"]
+                == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+            )
 
     @mock.patch.object(_agent_engines_utils, "_prepare")
     def test_update_agent_engine_config_full(self, mock_prepare):
@@ -1008,6 +1022,7 @@ def test_update_agent_engine_config_full(self, mock_prepare):
             extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
             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,
         )
         assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
         assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
@@ -1038,6 +1053,10 @@ def test_update_agent_engine_config_full(self, mock_prepare):
             config["spec"]["service_account"]
             == _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT
         )
+        assert (
+            config["spec"]["identity_type"]
+            == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+        )
         assert config["update_mask"] == ",".join(
             [
                 "display_name",
@@ -1048,8 +1067,28 @@ def test_update_agent_engine_config_full(self, mock_prepare):
                 "spec.class_methods",
                 "spec.deployment_spec.env",
                 "spec.deployment_spec.secret_env",
-                "spec.service_account",
                 "spec.agent_framework",
+                "spec.identity_type",
+                "spec.service_account",
+            ]
+        )
+
+    @mock.patch.object(_agent_engines_utils, "_prepare")
+    def test_update_agent_engine_clear_service_account(self, mock_prepare):
+        config = self.client.agent_engines._create_config(
+            mode="update",
+            service_account="",
+            identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
+        )
+        assert "service_account" not in config["spec"].keys()
+        assert (
+            config["spec"]["identity_type"]
+            == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+        )
+        assert config["update_mask"] == ",".join(
+            [
+                "spec.identity_type",
+                "spec.service_account",
             ]
         )
 
@@ -1488,6 +1527,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
                 env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT,
                 service_account=None,
+                identity_type=None,
                 context_spec=None,
                 psc_interface_config=None,
                 min_instances=None,
@@ -1543,6 +1583,7 @@ def test_create_agent_engine_with_custom_service_account(
                 },
                 "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1],
                 "service_account": _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
+                "identity_type": _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
                 "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK,
             },
         }
@@ -1564,6 +1605,7 @@ def test_create_agent_engine_with_custom_service_account(
                     extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
                     staging_bucket=_TEST_STAGING_BUCKET,
                     service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
+                    identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
                 ),
             )
             mock_create_config.assert_called_with(
@@ -1577,6 +1619,7 @@ def test_create_agent_engine_with_custom_service_account(
                 extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
                 env_vars=None,
                 service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
+                identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
                 context_spec=None,
                 psc_interface_config=None,
                 min_instances=None,
@@ -1608,6 +1651,7 @@ def test_create_agent_engine_with_custom_service_account(
                             "requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
                         },
                         "service_account": _TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
+                        "identity_type": _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
                     },
                 },
                 None,
@@ -1666,6 +1710,7 @@ def test_create_agent_engine_with_experimental_mode(
                 extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
                 env_vars=None,
                 service_account=None,
+                identity_type=None,
                 context_spec=None,
                 psc_interface_config=None,
                 min_instances=None,
@@ -1819,6 +1864,7 @@ def test_create_agent_engine_with_class_methods(
                 extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
                 env_vars=None,
                 service_account=None,
+                identity_type=None,
                 context_spec=None,
                 psc_interface_config=None,
                 min_instances=None,
@@ -1919,6 +1965,7 @@ def test_create_agent_engine_with_agent_framework(
                 entrypoint_object=None,
                 requirements_file=None,
                 agent_framework=_TEST_AGENT_FRAMEWORK,
+                identity_type=None,
             )
             request_mock.assert_called_with(
                 "post",

vertexai/_genai/agent_engines.py

@@ -907,6 +907,7 @@ def create(
         api_config = self._create_config(
             mode="create",
             agent=agent,
+            identity_type=config.identity_type,
             staging_bucket=config.staging_bucket,
             requirements=config.requirements,
             display_name=config.display_name,
@@ -971,6 +972,7 @@ def _create_config(
         *,
         mode: str,
         agent: Any = None,
+        identity_type: Optional[types.IdentityType] = None,
         staging_bucket: Optional[str] = None,
         requirements: Optional[Union[str, Sequence[str]]] = None,
         display_name: Optional[str] = None,
@@ -1189,9 +1191,6 @@ def _create_config(
                 )
                 update_masks.extend(deployment_update_masks)
                 agent_engine_spec["deployment_spec"] = deployment_spec
-            if service_account is not None:
-                agent_engine_spec["service_account"] = service_account
-                update_masks.append("spec.service_account")
 
             if agent_server_mode:
                 if not agent_engine_spec.get("deployment_spec"):
@@ -1209,6 +1208,21 @@ def _create_config(
                 )
             )
             update_masks.append("spec.agent_framework")
+
+        if identity_type is not None or service_account is not None:
+            if agent_engine_spec is None:
+                agent_engine_spec = {}
+
+            if identity_type is not None:
+                agent_engine_spec["identity_type"] = identity_type
+                update_masks.append("spec.identity_type")
+            if service_account is not None:
+                # Clear the field in case of empty service_account.
+                if service_account:
+                    agent_engine_spec["service_account"] = service_account
+                update_masks.append("spec.service_account")
+
+        if agent_engine_spec is not None:
             config["spec"] = agent_engine_spec
 
         if update_masks and mode == "update":
@@ -1414,6 +1428,7 @@ def update(
         api_config = self._create_config(
             mode="update",
             agent=agent,
+            identity_type=config.identity_type,
             staging_bucket=config.staging_bucket,
             requirements=config.requirements,
             display_name=config.display_name,

vertexai/_genai/types/__init__.py

@@ -430,6 +430,7 @@
 from .common import GetPromptConfig
 from .common import GetPromptConfigDict
 from .common import GetPromptConfigOrDict
+from .common import IdentityType
 from .common import Importance
 from .common import IntermediateExtractedMemory
 from .common import IntermediateExtractedMemoryDict
@@ -1798,6 +1799,7 @@
     "AcceleratorType",
     "Type",
     "JobState",
+    "IdentityType",
     "AgentServerMode",
     "ManagedTopicEnum",
     "Language",

vertexai/_genai/types/common.py

@@ -205,6 +205,17 @@ class JobState(_common.CaseInSensitiveEnum):
     """The job is partially succeeded, some results may be missing due to errors."""
 
 
+class IdentityType(_common.CaseInSensitiveEnum):
+    """The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used."""
+
+    IDENTITY_TYPE_UNSPECIFIED = "IDENTITY_TYPE_UNSPECIFIED"
+    """Default value. Use a custom service account if the `service_account` field is set, otherwise use the default Vertex AI Reasoning Engine Service Agent in the project. Same behavior as SERVICE_ACCOUNT."""
+    SERVICE_ACCOUNT = "SERVICE_ACCOUNT"
+    """Use a custom service account if the `service_account` field is set, otherwise use the default Vertex AI Reasoning Engine Service Agent in the project."""
+    AGENT_IDENTITY = "AGENT_IDENTITY"
+    """Use Agent Identity. The `service_account` field must not be set."""
+
+
 class AgentServerMode(_common.CaseInSensitiveEnum):
     """The agent server mode."""
 
@@ -4784,6 +4795,14 @@ class ReasoningEngineSpec(_common.BaseModel):
         default=None,
         description="""Optional. The specification of a Reasoning Engine deployment.""",
     )
+    effective_identity: Optional[str] = Field(
+        default=None,
+        description="""Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""",
+    )
+    identity_type: Optional[IdentityType] = Field(
+        default=None,
+        description="""Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""",
+    )
     package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field(
         default=None,
         description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""",
@@ -4810,6 +4829,12 @@ class ReasoningEngineSpecDict(TypedDict, total=False):
     deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict]
     """Optional. The specification of a Reasoning Engine deployment."""
 
+    effective_identity: Optional[str]
+    """Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)"""
+
+    identity_type: Optional[IdentityType]
+    """Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used."""
+
     package_spec: Optional[ReasoningEngineSpecPackageSpecDict]
     """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified."""
 
@@ -12905,6 +12930,9 @@ class AgentEngineConfig(_common.BaseModel):
 
       If not specified, the default Reasoning Engine P6SA service agent will be used.""",
     )
+    identity_type: Optional[IdentityType] = Field(
+        default=None, description="""The identity type to use for the Agent Engine."""
+    )
     context_spec: Optional[ReasoningEngineContextSpec] = Field(
         default=None,
         description="""The context spec to be used for the Agent Engine.""",
@@ -13057,6 +13085,9 @@ class AgentEngineConfigDict(TypedDict, total=False):
 
       If not specified, the default Reasoning Engine P6SA service agent will be used."""
 
+    identity_type: Optional[IdentityType]
+    """The identity type to use for the Agent Engine."""
+
     context_spec: Optional[ReasoningEngineContextSpecDict]
     """The context spec to be used for the Agent Engine."""