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."""
 
@@ -4790,6 +4800,14 @@ class ReasoningEngineSpec(_common.BaseModel):
         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.""",
     )
+    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`.""",
+    )
 
 
 class ReasoningEngineSpecDict(TypedDict, total=False):
@@ -4810,6 +4828,11 @@ class ReasoningEngineSpecDict(TypedDict, total=False):
     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."""
 
+    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`."""
 
 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