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

@@ -526,6 +526,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,
@@ -552,6 +555,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 = []
@@ -858,6 +862,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,
@@ -900,6 +905,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, "_prepare")
     def test_update_agent_engine_config_full(self, mock_prepare):
@@ -914,6 +923,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
@@ -944,21 +954,45 @@ 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",
                 "description",
+                "spec.identity_type",
+                "spec.service_account",
                 "spec.package_spec.pickle_object_gcs_uri",
                 "spec.package_spec.dependency_files_gcs_uri",
                 "spec.package_spec.requirements_gcs_uri",
                 "spec.deployment_spec.env",
                 "spec.deployment_spec.secret_env",
-                "spec.service_account",
                 "spec.class_methods",
                 "spec.agent_framework",
             ]
         )
 
+    @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",
+            ]
+        )
+
     def test_get_agent_operation(self):
         with mock.patch.object(
             self.client.agent_engines._api_client, "request"
@@ -1355,6 +1389,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,
@@ -1403,6 +1438,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,
             },
         }
@@ -1424,6 +1460,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(
@@ -1437,6 +1474,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,
@@ -1463,6 +1501,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,
@@ -1521,6 +1560,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,
@@ -1603,6 +1643,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,

vertexai/_genai/agent_engines.py

@@ -853,6 +853,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,
@@ -912,6 +913,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,
@@ -957,6 +959,17 @@ def _create_config(
         if labels is not None:
             update_masks.append("labels")
             config["labels"] = labels
+
+        agent_engine_spec: types.ReasoningEngineSpecDict = {}
+        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 is not None:
             project = self._api_client.project
             if project is None:
@@ -1013,9 +1026,7 @@ def _create_config(
                     gcs_dir_name,
                     _agent_engines_utils._REQUIREMENTS_FILE,
                 )
-            agent_engine_spec: types.ReasoningEngineSpecDict = {
-                "package_spec": package_spec,
-            }
+            agent_engine_spec["package_spec"] = package_spec
             if (
                 env_vars is not None
                 or psc_interface_config is not None
@@ -1037,9 +1048,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")
 
             update_masks.append("spec.class_methods")
             class_methods_spec = []
@@ -1076,7 +1084,10 @@ def _create_config(
                 _agent_engines_utils._get_agent_framework(agent=agent)
             )
             update_masks.append("spec.agent_framework")
+
+        if agent_engine_spec.items():
             config["spec"] = agent_engine_spec
+
         if update_masks and mode == "update":
             config["update_mask"] = ",".join(update_masks)
         return config
@@ -1280,6 +1291,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.py

@@ -360,6 +360,16 @@ class GenerateMemoriesResponseGeneratedMemoryAction(_common.CaseInSensitiveEnum)
     """The memory was deleted."""
 
 
+class IdentityType(_common.CaseInSensitiveEnum):
+    """The identity type for the agent engine."""
+    IDENTITY_TYPE_UNSPECIFIED = "IDENTITY_TYPE_UNSPECIFIED"
+    """The identity type is unspecified. 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."""
+    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 the Agent Identity to access the resources."""
+
+
 class SamplingConfig(_common.BaseModel):
     """Sampling config for a BigQuery request set."""
 
@@ -4788,7 +4798,15 @@ class ReasoningEngineSpec(_common.BaseModel):
     )
     service_account: Optional[str] = Field(
         default=None,
-        description="""Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used.""",
+        description="""Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used. To clear the service account, this field needs to be set to an empty string.""",
+    )
+    identity_type: Optional[IdentityType] = Field(
+        default=None,
+        description="""Optional. The identity type for the Reasoning Engine. If not specified, the default value is `IDENTITY_TYPE_UNSPECIFIED`.""",
+    )
+    effective_identity: Optional[str] = Field(
+        default=None,
+        description="""Output only. The identity to be used for the Reasoning Engine. If not specified, the default value is the service account specified in `service_account` or the Vertex AI Reasoning Engine Service Agent in the project if `service_account` is not specified.""",
     )
 
 
@@ -4808,8 +4826,13 @@ class ReasoningEngineSpecDict(TypedDict, total=False):
     """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."""
 
     service_account: Optional[str]
-    """Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used."""
+    """Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used. To clear the service account, this field needs to be set to an empty string."""
+
+    identity_type: Optional[IdentityType]
+    """Optional. The identity type for the Reasoning Engine. If not specified, the default value is `IDENTITY_TYPE_UNSPECIFIED`."""
 
+    effective_identity: Optional[str]
+    """Output only. The identity to be used for the Reasoning Engine. If not specified, the default value is the service account specified in `service_account` or the Vertex AI Reasoning Engine Service Agent in the project if `service_account` is not specified."""
 
 ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict]
 
@@ -12198,6 +12221,10 @@ class AgentEngineConfig(_common.BaseModel):
     agent_server_mode: Optional[AgentServerMode] = Field(
         default=None, description="""The agent server mode to use for deployment."""
     )
+    identity_type: Optional[IdentityType] = Field(
+        default=None,
+        description="""The identity type to use for the Agent Engine.""",
+    )
     class_methods: Optional[list[dict[str, Any]]] = Field(
         default=None,
         description="""The class methods to be used for the Agent Engine.
@@ -12293,6 +12320,9 @@ class AgentEngineConfigDict(TypedDict, total=False):
     agent_server_mode: Optional[AgentServerMode]
     """The agent server mode to use for deployment."""
 
+    identity_type: Optional[IdentityType]
+    """The identity type to use for the Agent Engine."""
+
     class_methods: Optional[list[dict[str, Any]]]
     """The class methods to be used for the Agent Engine.
       If specified, they'll override the class methods that are autogenerated by