diff --git a/tests/unit/vertex_langchain/test_agent_engines.py b/tests/unit/vertex_langchain/test_agent_engines.py index 8c3610577e..e12df77471 100644 --- a/tests/unit/vertex_langchain/test_agent_engines.py +++ b/tests/unit/vertex_langchain/test_agent_engines.py @@ -613,6 +613,10 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_INSTALLATION_SCRIPT_PATH = f"{_TEST_INSTALLATION_SUBDIR}/install_package.sh" _TEST_CUSTOM_SERVICE_ACCOUNT = "test-custom-service-account" +_TEST_AGENT_IDENTITY_TYPE = types.ReasoningEngineSpec.IdentityType.AGENT_IDENTITY +_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE = ( + types.ReasoningEngineSpec.IdentityType.SERVICE_ACCOUNT +) _TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG = { "network_attachment": "test-network-attachment", "dns_peering_configs": [ @@ -1284,6 +1288,72 @@ def test_create_agent_engine_with_service_account( retry=_TEST_RETRY, ) + def test_create_agent_engine_with_agent_identity( + self, + create_agent_engine_mock, + cloud_storage_create_bucket_mock, + tarfile_open_mock, + cloudpickle_dump_mock, + cloudpickle_load_mock, + importlib_metadata_version_mock, + get_agent_engine_mock, + get_gca_resource_mock, + ): + agent_engines.create( + self.test_agent, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + test_spec = types.ReasoningEngineSpec( + package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, + agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA) + create_agent_engine_mock.assert_called_with( + parent=_TEST_PARENT, + reasoning_engine=types.ReasoningEngine( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + spec=test_spec, + ), + ) + + def test_create_agent_engine_with_service_account_identity_type( + self, + create_agent_engine_mock, + cloud_storage_create_bucket_mock, + tarfile_open_mock, + cloudpickle_dump_mock, + cloudpickle_load_mock, + importlib_metadata_version_mock, + get_agent_engine_mock, + get_gca_resource_mock, + ): + agent_engines.create( + self.test_agent, + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + ) + test_spec = types.ReasoningEngineSpec( + package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, + agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + ) + test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA) + create_agent_engine_mock.assert_called_with( + parent=_TEST_PARENT, + reasoning_engine=types.ReasoningEngine( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + spec=test_spec, + ), + ) + def test_create_agent_engine_with_psc_interface_config( self, create_agent_engine_mock, @@ -1809,6 +1879,45 @@ def test_get_agent_framework( ), ), ), + ( + "Update the agent_engine with identity_type attribute", + {"identity_type": _TEST_AGENT_IDENTITY_TYPE}, + types.reasoning_engine_service.UpdateReasoningEngineRequest( + reasoning_engine=types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=types.ReasoningEngineSpec( + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ), + ), + update_mask=field_mask_pb2.FieldMask( + paths=[ + "spec.identity_type", + ], + ), + ), + ), + ( + "Update the agent_engine with service_account and identity_type", + { + "service_account": _TEST_CUSTOM_SERVICE_ACCOUNT, + "identity_type": _TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + }, + types.reasoning_engine_service.UpdateReasoningEngineRequest( + reasoning_engine=types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=types.ReasoningEngineSpec( + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + ), + ), + update_mask=field_mask_pb2.FieldMask( + paths=[ + "spec.service_account", + "spec.identity_type", + ], + ), + ), + ), ( "Update the agent_engine with psc_interface_config attribute", {"psc_interface_config": _TEST_AGENT_ENGINE_PSC_INTERFACE_CONFIG}, @@ -3406,14 +3515,45 @@ def test_update_agent_engine_with_no_updates( "At least one of `agent_engine`, `requirements`, " "`extra_packages`, `display_name`, `description`, " "`env_vars`, `build_options`, `service_account`, " - "`psc_interface_config`, `min_instances`, `max_instances`, " - "`resource_limits`, `container_concurrency`, or " - "`encryption_spec` must be specified." + "`identity_type`, `psc_interface_config`, `min_instances`, " + "`max_instances`, `resource_limits`, `container_concurrency`, " + "or `encryption_spec` must be specified." ), ): test_agent_engine = _generate_agent_engine_to_update() test_agent_engine.update() + def test_create_agent_engine_with_agent_identity_and_service_account(self): + with pytest.raises( + ValueError, + match=( + "`service_account` must not be specified when `identity_type` " + "is `AGENT_IDENTITY`" + ), + ): + agent_engines.create( + CapitalizeEngine(), + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + + def test_update_agent_engine_with_agent_identity_and_service_account( + self, + update_agent_engine_mock, + ): + with pytest.raises( + ValueError, + match=( + "`service_account` must not be specified when `identity_type` " + "is `AGENT_IDENTITY`" + ), + ): + test_agent_engine = _generate_agent_engine_to_update() + test_agent_engine.update( + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + def test_create_class_methods_spec_with_registered_operation_not_found(self): with pytest.raises( ValueError, diff --git a/tests/unit/vertex_langchain/test_reasoning_engines.py b/tests/unit/vertex_langchain/test_reasoning_engines.py index 019dc214d9..45cd030b93 100644 --- a/tests/unit/vertex_langchain/test_reasoning_engines.py +++ b/tests/unit/vertex_langchain/test_reasoning_engines.py @@ -348,6 +348,42 @@ class ListClass: _TEST_INPUT_REASONING_ENGINE_OBJ.spec.class_methods.append( _TEST_REASONING_ENGINE_QUERY_SCHEMA ) +_TEST_CUSTOM_SERVICE_ACCOUNT = "test-custom-service-account" +_TEST_AGENT_IDENTITY_TYPE = types.ReasoningEngineSpec.IdentityType.AGENT_IDENTITY +_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE = ( + types.ReasoningEngineSpec.IdentityType.SERVICE_ACCOUNT +) +_TEST_INPUT_REASONING_ENGINE_OBJ_WITH_SERVICE_ACCOUNT = types.ReasoningEngine( + display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME, + spec=types.ReasoningEngineSpec( + package_spec=types.ReasoningEngineSpec.PackageSpec( + python_version=f"{sys.version_info.major}.{sys.version_info.minor}", + pickle_object_gcs_uri=_TEST_REASONING_ENGINE_GCS_URI, + dependency_files_gcs_uri=_TEST_REASONING_ENGINE_DEPENDENCY_FILES_GCS_URI, + requirements_gcs_uri=_TEST_REASONING_ENGINE_REQUIREMENTS_GCS_URI, + ), + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + ), +) +_TEST_INPUT_REASONING_ENGINE_OBJ_WITH_SERVICE_ACCOUNT.spec.class_methods.append( + _TEST_REASONING_ENGINE_QUERY_SCHEMA +) +_TEST_INPUT_REASONING_ENGINE_OBJ_WITH_AGENT_IDENTITY = types.ReasoningEngine( + display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME, + spec=types.ReasoningEngineSpec( + package_spec=types.ReasoningEngineSpec.PackageSpec( + python_version=f"{sys.version_info.major}.{sys.version_info.minor}", + pickle_object_gcs_uri=_TEST_REASONING_ENGINE_GCS_URI, + dependency_files_gcs_uri=_TEST_REASONING_ENGINE_DEPENDENCY_FILES_GCS_URI, + requirements_gcs_uri=_TEST_REASONING_ENGINE_REQUIREMENTS_GCS_URI, + ), + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ), +) +_TEST_INPUT_REASONING_ENGINE_OBJ_WITH_AGENT_IDENTITY.spec.class_methods.append( + _TEST_REASONING_ENGINE_QUERY_SCHEMA +) _TEST_REASONING_ENGINE_OBJ = types.ReasoningEngine( name=_TEST_REASONING_ENGINE_RESOURCE_NAME, display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME, @@ -794,6 +830,49 @@ def test_create_reasoning_engine( retry=_TEST_RETRY, ) + def test_create_reasoning_engine_with_service_account( + self, + create_reasoning_engine_mock, + cloud_storage_create_bucket_mock, + tarfile_open_mock, + cloudpickle_dump_mock, + get_reasoning_engine_mock, + get_gca_resource_mock, + ): + reasoning_engines.ReasoningEngine.create( + self.test_app, + display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME, + requirements=_TEST_REASONING_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_REASONING_ENGINE_EXTRA_PACKAGE_PATH], + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + ) + create_reasoning_engine_mock.assert_called_with( + parent=_TEST_PARENT, + reasoning_engine=_TEST_INPUT_REASONING_ENGINE_OBJ_WITH_SERVICE_ACCOUNT, + ) + + def test_create_reasoning_engine_with_agent_identity( + self, + create_reasoning_engine_mock, + cloud_storage_create_bucket_mock, + tarfile_open_mock, + cloudpickle_dump_mock, + get_reasoning_engine_mock, + get_gca_resource_mock, + ): + reasoning_engines.ReasoningEngine.create( + self.test_app, + display_name=_TEST_REASONING_ENGINE_DISPLAY_NAME, + requirements=_TEST_REASONING_ENGINE_REQUIREMENTS, + extra_packages=[_TEST_REASONING_ENGINE_EXTRA_PACKAGE_PATH], + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + create_reasoning_engine_mock.assert_called_with( + parent=_TEST_PARENT, + reasoning_engine=_TEST_INPUT_REASONING_ENGINE_OBJ_WITH_AGENT_IDENTITY, + ) + @pytest.mark.usefixtures("caplog") def test_create_reasoning_engine_warn_resource_name( self, @@ -990,6 +1069,56 @@ def test_create_reasoning_engine_requirements_from_file( update_mask=field_mask_pb2.FieldMask(paths=["description"]), ), ), + ( + "Update the service_account", + {"service_account": _TEST_CUSTOM_SERVICE_ACCOUNT}, + types.reasoning_engine_service.UpdateReasoningEngineRequest( + reasoning_engine=types.ReasoningEngine( + name=_TEST_REASONING_ENGINE_RESOURCE_NAME, + spec=types.ReasoningEngineSpec( + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + ), + ), + update_mask=field_mask_pb2.FieldMask( + paths=["spec.service_account"] + ), + ), + ), + ( + "Update the identity_type", + {"identity_type": _TEST_AGENT_IDENTITY_TYPE}, + types.reasoning_engine_service.UpdateReasoningEngineRequest( + reasoning_engine=types.ReasoningEngine( + name=_TEST_REASONING_ENGINE_RESOURCE_NAME, + spec=types.ReasoningEngineSpec( + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ), + ), + update_mask=field_mask_pb2.FieldMask(paths=["spec.identity_type"]), + ), + ), + ( + "Update the service_account and identity_type", + { + "service_account": _TEST_CUSTOM_SERVICE_ACCOUNT, + "identity_type": _TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + }, + types.reasoning_engine_service.UpdateReasoningEngineRequest( + reasoning_engine=types.ReasoningEngine( + name=_TEST_REASONING_ENGINE_RESOURCE_NAME, + spec=types.ReasoningEngineSpec( + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_SERVICE_ACCOUNT_IDENTITY_TYPE, + ), + ), + update_mask=field_mask_pb2.FieldMask( + paths=[ + "spec.service_account", + "spec.identity_type", + ] + ), + ), + ), ], ) def test_update_reasoning_engine( @@ -2010,13 +2139,44 @@ def test_update_reasoning_engine_with_no_updates( ValueError, match=( "At least one of `reasoning_engine`, `requirements`, " - "`extra_packages`, `display_name`, or `description` " - "must be specified." + "`extra_packages`, `display_name`, `description`, " + "`service_account`, or `identity_type` must be specified." ), ): test_reasoning_engine = _generate_reasoning_engine_to_update() test_reasoning_engine.update() + def test_create_reasoning_engine_with_agent_identity_and_service_account(self): + with pytest.raises( + ValueError, + match=( + "`service_account` must not be specified when `identity_type` " + "is `AGENT_IDENTITY`" + ), + ): + reasoning_engines.ReasoningEngine.create( + CapitalizeEngine(), + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + + def test_update_reasoning_engine_with_agent_identity_and_service_account( + self, + update_reasoning_engine_mock, + ): + with pytest.raises( + ValueError, + match=( + "`service_account` must not be specified when `identity_type` " + "is `AGENT_IDENTITY`" + ), + ): + test_reasoning_engine = _generate_reasoning_engine_to_update() + test_reasoning_engine.update( + service_account=_TEST_CUSTOM_SERVICE_ACCOUNT, + identity_type=_TEST_AGENT_IDENTITY_TYPE, + ) + def test_create_class_methods_spec_with_registered_operation_not_found(self): with pytest.raises( ValueError, diff --git a/vertexai/agent_engines/__init__.py b/vertexai/agent_engines/__init__.py index 3d183ef345..feb333e792 100644 --- a/vertexai/agent_engines/__init__.py +++ b/vertexai/agent_engines/__init__.py @@ -76,6 +76,7 @@ def create( ] = None, build_options: Optional[Dict[str, Sequence[str]]] = None, service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, @@ -159,6 +160,14 @@ def create( Optional. The service account to be used for the Agent Engine. If not specified, the default reasoning engine service agent service account will be used. + identity_type (ReasoningEngineSpec.IdentityType): + Optional. The identity type to be used by the Agent Engine at + runtime. Set it to `AGENT_IDENTITY` to run the Agent Engine under a + resource-scoped workload identity with no default project + permissions, instead of the shared Reasoning Engine Service Agent. + `service_account` must not be specified when `AGENT_IDENTITY` is + used. If not specified, `service_account` is used if set, otherwise + the default Reasoning Engine Service Agent is used. psc_interface_config (PscInterfaceConfig): Optional. The PSC interface config for the Agent Engine. If not specified, the default PSC interface config will be used. @@ -186,6 +195,8 @@ def create( ValueError: If the `location` was not set using `vertexai.init`. ValueError: If the `staging_bucket` was not set using vertexai.init. ValueError: If the `staging_bucket` does not start with "gs://". + ValueError: If `service_account` is specified and `identity_type` is + `AGENT_IDENTITY`. FileNotFoundError: If `extra_packages` includes a file or directory that does not exist. IOError: If requirements is a string that corresponds to a @@ -201,6 +212,7 @@ def create( env_vars=env_vars, build_options=build_options, service_account=service_account, + identity_type=identity_type, psc_interface_config=psc_interface_config, min_instances=min_instances, max_instances=max_instances, @@ -293,6 +305,7 @@ def update( ] = None, build_options: Optional[Dict[str, Sequence[str]]] = None, service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, @@ -349,6 +362,13 @@ def update( Optional. The service account to be used for the Agent Engine. If not specified, the default reasoning engine service agent service account will be used. + identity_type (ReasoningEngineSpec.IdentityType): + Optional. The identity type to be used by the Agent Engine at + runtime. Set it to `AGENT_IDENTITY` to run the Agent Engine under a + resource-scoped workload identity with no default project + permissions, instead of the shared Reasoning Engine Service Agent. + `service_account` must not be specified when `AGENT_IDENTITY` is + used. If not specified, the existing identity type will be used. min_instances (int): Optional. The minimum number of instances to run the Agent Engine. If not specified, the default value will be used. @@ -376,6 +396,8 @@ def update( ValueError: if none of `display_name`, `description`, `requirements`, `extra_packages`, `agent_engine`, or `build_options` were specified. + ValueError: If `service_account` is specified and `identity_type` is + `AGENT_IDENTITY`. IOError: If requirements is a string that corresponds to a nonexistent file. """ @@ -390,6 +412,7 @@ def update( env_vars=env_vars, build_options=build_options, service_account=service_account, + identity_type=identity_type, psc_interface_config=psc_interface_config, min_instances=min_instances, max_instances=max_instances, diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index 5a87dfebd1..6e9ae0f18c 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -387,6 +387,7 @@ def create( ] = None, build_options: Optional[Dict[str, Sequence[str]]] = None, service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, @@ -475,6 +476,15 @@ def create( Optional. The service account to be used for the Agent Engine. If not specified, the default reasoning engine service agent service account will be used. + identity_type (aip_types.ReasoningEngineSpec.IdentityType): + Optional. The identity type to be used by the Agent Engine at + runtime. Set it to `AGENT_IDENTITY` to run the Agent Engine + under a resource-scoped workload identity with no default + project permissions, instead of the shared Reasoning Engine + Service Agent. `service_account` must not be specified when + `AGENT_IDENTITY` is used. If not specified, `service_account` + is used if set, otherwise the default Reasoning Engine Service + Agent is used. psc_interface_config (aip_types.PscInterfaceConfig): Optional. The Private Service Connect interface config for the Agent Engine. @@ -505,6 +515,8 @@ def create( ValueError: If the `staging_bucket` does not start with "gs://". ValueError: If `extra_packages` is specified but `agent_engine` is None. ValueError: If `requirements` is specified but `agent_engine` is None. + ValueError: If `service_account` is specified and `identity_type` is + `AGENT_IDENTITY`. ValueError: If `env_vars` has a dictionary entry that does not correspond to a SecretRef. ValueError: If `env_vars` is a list which contains a string that @@ -518,6 +530,10 @@ def create( """ sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" _validate_sys_version_or_raise(sys_version) + _validate_identity_type_or_raise( + identity_type=identity_type, + service_account=service_account, + ) gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME staging_bucket = initializer.global_config.staging_bucket @@ -611,6 +627,8 @@ def create( agent_engine_spec.class_methods.extend(class_methods_spec) if service_account: agent_engine_spec.service_account = service_account + if identity_type is not None: + agent_engine_spec.identity_type = identity_type reasoning_engine.spec = agent_engine_spec reasoning_engine.spec.agent_framework = _get_agent_framework(agent_engine) operation_future = sdk_resource.api_client.create_reasoning_engine( @@ -666,6 +684,7 @@ def update( ] = None, build_options: Optional[Dict[str, Sequence[str]]] = None, service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, @@ -725,6 +744,14 @@ def update( Optional. The service account to be used for the Agent Engine. If not specified, the default reasoning engine service agent service account will be used. + identity_type (aip_types.ReasoningEngineSpec.IdentityType): + Optional. The identity type to be used by the Agent Engine at + runtime. Set it to `AGENT_IDENTITY` to run the Agent Engine + under a resource-scoped workload identity with no default + project permissions, instead of the shared Reasoning Engine + Service Agent. `service_account` must not be specified when + `AGENT_IDENTITY` is used. If not specified, the existing + identity type will be used. psc_interface_config (aip_types.PscInterfaceConfig): Optional. The Private Service Connect interface config for the Agent Engine. @@ -761,6 +788,8 @@ def update( that does not exist. ValueError: if none of `display_name`, `description`, `requirements`, `extra_packages`, `env_vars`, or `agent_engine` were specified. + ValueError: If `service_account` is specified and `identity_type` is + `AGENT_IDENTITY`. IOError: If requirements is a string that corresponds to a nonexistent file. """ @@ -780,6 +809,7 @@ def update( env_vars, build_options, service_account, + identity_type is not None, psc_interface_config, min_instances is not None, max_instances is not None, @@ -792,10 +822,14 @@ def update( "At least one of `agent_engine`, `requirements`, " "`extra_packages`, `display_name`, `description`, " "`env_vars`, `build_options`, `service_account`, " - "`psc_interface_config`, `min_instances`, `max_instances`, " - "`resource_limits`, `container_concurrency`, or " - "`encryption_spec` must be specified." + "`identity_type`, `psc_interface_config`, `min_instances`, " + "`max_instances`, `resource_limits`, `container_concurrency`, " + "or `encryption_spec` must be specified." ) + _validate_identity_type_or_raise( + identity_type=identity_type, + service_account=service_account, + ) if requirements is not None: requirements = _validate_requirements_or_raise( agent_engine=agent_engine, @@ -835,6 +869,7 @@ def update( description=description, env_vars=env_vars, service_account=service_account, + identity_type=identity_type, psc_interface_config=psc_interface_config, min_instances=min_instances, max_instances=max_instances, @@ -940,6 +975,31 @@ def _validate_staging_bucket_or_raise(staging_bucket: Optional[str]) -> str: return staging_bucket +def _validate_identity_type_or_raise( + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType], + service_account: Optional[str], +) -> None: + """Tries to validate the identity type against the service account. + + Args: + identity_type: The identity type to be used by the Agent Engine. + service_account: The service account to be used by the Agent Engine. + + Raises: + ValueError: If `identity_type` is `AGENT_IDENTITY` and `service_account` + is also specified, because the two are mutually exclusive. + """ + if ( + identity_type == aip_types.ReasoningEngineSpec.IdentityType.AGENT_IDENTITY + and service_account + ): + raise ValueError( + "`service_account` must not be specified when `identity_type` is " + "`AGENT_IDENTITY`, because the Agent Engine runs under a " + "resource-scoped workload identity instead of a service account." + ) + + def _validate_agent_engine_or_raise( agent_engine: _AgentEngineInterface, logger: base.Logger = _LOGGER, @@ -1451,6 +1511,7 @@ def _generate_update_request_or_raise( Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] ] = None, service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, min_instances: Optional[int] = None, max_instances: Optional[int] = None, @@ -1518,6 +1579,10 @@ def _generate_update_request_or_raise( is_spec_update = True update_masks.append("spec.service_account") agent_engine_spec.service_account = service_account + if identity_type is not None: + is_spec_update = True + update_masks.append("spec.identity_type") + agent_engine_spec.identity_type = identity_type agent_engine_message = aip_types.ReasoningEngine(name=resource_name) if is_spec_update: @@ -1536,8 +1601,8 @@ def _generate_update_request_or_raise( if not update_masks: raise ValueError( "At least one of `agent_engine`, `requirements`, `extra_packages`, " - "`display_name`, `description`, `env_vars`, or " - "`encryption_spec` must be specified." + "`display_name`, `description`, `env_vars`, `service_account`, " + "`identity_type`, or `encryption_spec` must be specified." ) return reasoning_engine_service.UpdateReasoningEngineRequest( reasoning_engine=agent_engine_message, diff --git a/vertexai/reasoning_engines/_reasoning_engines.py b/vertexai/reasoning_engines/_reasoning_engines.py index 7d94cda0bc..ac100c6141 100644 --- a/vertexai/reasoning_engines/_reasoning_engines.py +++ b/vertexai/reasoning_engines/_reasoning_engines.py @@ -161,6 +161,8 @@ def create( gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, sys_version: Optional[str] = None, extra_packages: Optional[Sequence[str]] = None, + service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, ) -> "ReasoningEngine": """Creates a new ReasoningEngine. @@ -229,6 +231,19 @@ def create( sys.version_info. extra_packages (Sequence[str]): Optional. The set of extra user-provided packages (if any). + service_account (str): + Optional. The service account that the Reasoning Engine runs + as. If not specified, the default Reasoning Engine Service + Agent in the project will be used. + identity_type (aip_types.ReasoningEngineSpec.IdentityType): + Optional. The identity type to be used by the Reasoning Engine + at runtime. Set it to `AGENT_IDENTITY` to run the Reasoning + Engine under a resource-scoped workload identity with no + default project permissions, instead of the shared Reasoning + Engine Service Agent. `service_account` must not be specified + when `AGENT_IDENTITY` is used. If not specified, + `service_account` is used if set, otherwise the default + Reasoning Engine Service Agent is used. Returns: ReasoningEngine: The Reasoning Engine that was created. @@ -239,6 +254,8 @@ def create( ValueError: If the `location` was not set using `vertexai.init`. ValueError: If the `staging_bucket` was not set using vertexai.init. ValueError: If the `staging_bucket` does not start with "gs://". + ValueError: If `service_account` is specified and `identity_type` + is `AGENT_IDENTITY`. FileNotFoundError: If `extra_packages` includes a file or directory that does not exist. IOError: If requirements is a string that corresponds to a @@ -247,6 +264,10 @@ def create( if not sys_version: sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" _validate_sys_version_or_raise(sys_version) + _validate_identity_type_or_raise( + identity_type=identity_type, + service_account=service_account, + ) reasoning_engine = _validate_reasoning_engine_or_raise(reasoning_engine) requirements = _validate_requirements_or_raise(requirements) extra_packages = _validate_extra_packages_or_raise(extra_packages) @@ -305,6 +326,10 @@ def create( reasoning_engine, _get_registered_operations(reasoning_engine) ) reasoning_engine_spec.class_methods.extend(class_methods_spec) + if service_account: + reasoning_engine_spec.service_account = service_account + if identity_type is not None: + reasoning_engine_spec.identity_type = identity_type operation_future = sdk_resource.api_client.create_reasoning_engine( parent=initializer.global_config.common_location_path( project=sdk_resource.project, location=sdk_resource.location @@ -351,6 +376,8 @@ def update( gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, sys_version: Optional[str] = None, extra_packages: Optional[Sequence[str]] = None, + service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, ) -> "ReasoningEngine": """Updates an existing ReasoningEngine. @@ -390,6 +417,18 @@ def update( it is not specified, the existing extra packages will be used. If it is set to an empty list, the existing extra packages will be removed. + service_account (str): + Optional. The service account that the Reasoning Engine runs + as. If not specified, the existing service account will be + used. + identity_type (aip_types.ReasoningEngineSpec.IdentityType): + Optional. The identity type to be used by the Reasoning Engine + at runtime. Set it to `AGENT_IDENTITY` to run the Reasoning + Engine under a resource-scoped workload identity with no + default project permissions, instead of the shared Reasoning + Engine Service Agent. `service_account` must not be specified + when `AGENT_IDENTITY` is used. If not specified, the existing + identity type will be used. Returns: ReasoningEngine: The Reasoning Engine that was updated. @@ -403,6 +442,8 @@ def update( ValueError: if none of `display_name`, `description`, `requirements`, `extra_packages`, or `reasoning_engine` were specified. + ValueError: If `service_account` is specified and `identity_type` + is `AGENT_IDENTITY`. IOError: If requirements is a string that corresponds to a nonexistent file. """ @@ -418,13 +459,19 @@ def update( extra_packages, display_name, description, + service_account, + identity_type is not None, ] ): raise ValueError( "At least one of `reasoning_engine`, `requirements`, " - "`extra_packages`, `display_name`, or `description` must be " - "specified." + "`extra_packages`, `display_name`, `description`, " + "`service_account`, or `identity_type` must be specified." ) + _validate_identity_type_or_raise( + identity_type=identity_type, + service_account=service_account, + ) if sys_version: _LOGGER.warning("Updated sys_version is not supported.") if requirements is not None: @@ -456,6 +503,8 @@ def update( extra_packages=extra_packages, display_name=display_name, description=description, + service_account=service_account, + identity_type=identity_type, ) operation_future = self.api_client.update_reasoning_engine( request=update_request @@ -523,6 +572,32 @@ def _validate_staging_bucket_or_raise(staging_bucket: str) -> str: raise ValueError(f"{staging_bucket=} must start with `gs://`") +def _validate_identity_type_or_raise( + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType], + service_account: Optional[str], +) -> None: + """Tries to validate the identity type against the service account. + + Args: + identity_type: The identity type to be used by the Reasoning Engine. + service_account: The service account to be used by the Reasoning + Engine. + + Raises: + ValueError: If `identity_type` is `AGENT_IDENTITY` and `service_account` + is also specified, because the two are mutually exclusive. + """ + if ( + identity_type == aip_types.ReasoningEngineSpec.IdentityType.AGENT_IDENTITY + and service_account + ): + raise ValueError( + "`service_account` must not be specified when `identity_type` is " + "`AGENT_IDENTITY`, because the Reasoning Engine runs under a " + "resource-scoped workload identity instead of a service account." + ) + + def _validate_reasoning_engine_or_raise( reasoning_engine: Union[Queryable, OperationRegistrable, StreamQueryable], ) -> Union[Queryable, OperationRegistrable, StreamQueryable]: @@ -723,6 +798,8 @@ def _generate_update_request_or_raise( extra_packages: Optional[Sequence[str]] = None, display_name: Optional[str] = None, description: Optional[str] = None, + service_account: Optional[str] = None, + identity_type: Optional[aip_types.ReasoningEngineSpec.IdentityType] = None, ) -> reasoning_engine_service.UpdateReasoningEngineRequest: """Tries to generates the update request for the reasoning engine.""" is_spec_update = False @@ -758,10 +835,19 @@ def _generate_update_request_or_raise( ) reasoning_engine_spec.class_methods.extend(class_methods_spec) update_masks.append("spec.class_methods") + if service_account is not None: + is_spec_update = True + update_masks.append("spec.service_account") + reasoning_engine_spec.service_account = service_account + if identity_type is not None: + is_spec_update = True + update_masks.append("spec.identity_type") + reasoning_engine_spec.identity_type = identity_type reasoning_engine_message = aip_types.ReasoningEngine(name=resource_name) if is_spec_update: - reasoning_engine_spec.package_spec = package_spec + if package_spec: + reasoning_engine_spec.package_spec = package_spec reasoning_engine_message.spec = reasoning_engine_spec if display_name: reasoning_engine_message.display_name = display_name @@ -772,8 +858,8 @@ def _generate_update_request_or_raise( if not update_masks: raise ValueError( "At least one of `reasoning_engine`, `requirements`, " - "`extra_packages`, `display_name`, or `description` must be " - "specified." + "`extra_packages`, `display_name`, `description`, " + "`service_account`, or `identity_type` must be specified." ) return reasoning_engine_service.UpdateReasoningEngineRequest( reasoning_engine=reasoning_engine_message, @@ -897,8 +983,7 @@ def _register_api_methods_or_raise(obj: "ReasoningEngine"): api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA) if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema: raise ValueError( - f"Operation schema {operation_schema} does not" - " contain a `name` field." + f"Operation schema {operation_schema} does not contain a `name` field." ) method_name = operation_schema.get(_METHOD_NAME_KEY_IN_SCHEMA) method_description = operation_schema.get("description")