feat: Add encryption_spec support to Agent Engine genai sdk.

PiperOrigin-RevId: 797986621

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -501,6 +501,7 @@ def register_operations(self) -> Dict[str, List[str]]:
  "memory": "4Gi",
 }
 _TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY = 4
+_TEST_AGENT_ENGINE_ENCRYPTION_SPEC = {"kms_key_name": "test-kms-key"}
 _TEST_AGENT_ENGINE_SPEC = _genai_types.ReasoningEngineSpecDict(
  agent_framework=_TEST_AGENT_ENGINE_FRAMEWORK,
  class_methods=[_TEST_AGENT_ENGINE_CLASS_METHOD_1],
@@ -803,6 +804,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
  max_instances=_TEST_AGENT_ENGINE_MAX_INSTANCES,
  resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS,
  container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
+ encryption_spec=_TEST_AGENT_ENGINE_ENCRYPTION_SPEC,
  )
  assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
  assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
@@ -833,6 +835,7 @@ def test_create_agent_engine_config_full(self, mock_prepare):
  "resource_limits": _TEST_AGENT_ENGINE_RESOURCE_LIMITS,
  "container_concurrency": _TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
  }
+ assert config["encryption_spec"] == _TEST_AGENT_ENGINE_ENCRYPTION_SPEC
  assert config["spec"]["class_methods"] == [_TEST_AGENT_ENGINE_CLASS_METHOD_1]
 
  @mock.patch.object(_agent_engines_utils, "_prepare")
@@ -1288,6 +1291,7 @@ def test_create_agent_engine_with_env_vars_dict(
  max_instances=None,
  resource_limits=None,
  container_concurrency=None,
+ encryption_spec=None,
  )
  request_mock.assert_called_with(
  "post",

vertexai/_genai/agent_engines.py

@@ -108,6 +108,17 @@ def _PscInterfaceConfig_to_vertex(
  return to_object
 
 
+def _EncryptionSpec_to_vertex(
+ from_object: Union[dict[str, Any], object],
+ parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+ to_object: dict[str, Any] = {}
+ if getv(from_object, ["kms_key_name"]) is not None:
+ setv(to_object, ["kmsKeyName"], getv(from_object, ["kms_key_name"]))
+
+ return to_object
+
+
 def _CreateAgentEngineConfig_to_vertex(
  from_object: Union[dict[str, Any], object],
  parent_object: Optional[dict[str, Any]] = None,
@@ -145,6 +156,15 @@ def _CreateAgentEngineConfig_to_vertex(
  ),
  )
 
+ if getv(from_object, ["encryption_spec"]) is not None:
+ setv(
+ parent_object,
+ ["encryptionSpec"],
+ _EncryptionSpec_to_vertex(
+ getv(from_object, ["encryption_spec"]), to_object
+ ),
+ )
+
  return to_object
 
 
@@ -332,6 +352,15 @@ def _UpdateAgentEngineConfig_to_vertex(
  ),
  )
 
+ if getv(from_object, ["encryption_spec"]) is not None:
+ setv(
+ parent_object,
+ ["encryptionSpec"],
+ _EncryptionSpec_to_vertex(
+ getv(from_object, ["encryption_spec"]), to_object
+ ),
+ )
+
  if getv(from_object, ["update_mask"]) is not None:
  setv(
  parent_object,
@@ -1065,6 +1094,7 @@ def create(
  max_instances=config.max_instances,
  resource_limits=config.resource_limits,
  container_concurrency=config.container_concurrency,
+ encryption_spec=config.encryption_spec,
  )
  operation = self._create(config=api_config)
  # TODO: Use a more specific link.
@@ -1118,6 +1148,7 @@ def _create_config(
  max_instances: Optional[int] = None,
  resource_limits: Optional[dict[str, str]] = None,
  container_concurrency: Optional[int] = None,
+ encryption_spec: Optional[types.EncryptionSpecDict] = None,
  ) -> types.UpdateAgentEngineConfigDict:
  import sys
 
@@ -1137,7 +1168,11 @@ def _create_config(
  update_masks.append("description")
  config["description"] = description
  if context_spec is not None:
+ update_masks.append("context_spec")
  config["context_spec"] = context_spec
+ if encryption_spec is not None:
+ update_masks.append("encryption_spec")
+ config["encryption_spec"] = encryption_spec
  if agent is not None:
  project = self._api_client.project
  if project is None:

vertexai/_genai/types.py

@@ -3170,7 +3170,7 @@ class CustomJobSpecDict(TypedDict, total=False):
 
 
 class EncryptionSpec(_common.BaseModel):
- """Represents a customer-managed encryption key spec that can be applied to a top-level resource."""
+ """The encryption spec."""
 
  kms_key_name: Optional[str] = Field(
  default=None,
@@ -3179,7 +3179,7 @@ class EncryptionSpec(_common.BaseModel):
 
 
 class EncryptionSpecDict(TypedDict, total=False):
- """Represents a customer-managed encryption key spec that can be applied to a top-level resource."""
+ """The encryption spec."""
 
  kms_key_name: Optional[str]
  """Required. The Cloud KMS resource identifier of the customer managed encryption key used to protect a resource. Has the form: `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. The key needs to be in the same region as where the compute resource is created."""
@@ -4029,6 +4029,10 @@ class CreateAgentEngineConfig(_common.BaseModel):
  Recommended value: 2 * cpu + 1. Defaults to 9.
  """,
  )
+ encryption_spec: Optional[EncryptionSpec] = Field(
+ default=None,
+ description="""The encryption spec to be used for the Agent Engine.""",
+ )
 
 
 class CreateAgentEngineConfigDict(TypedDict, total=False):
@@ -4080,6 +4084,9 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
  Recommended value: 2 * cpu + 1. Defaults to 9.
  """
 
+ encryption_spec: Optional[EncryptionSpecDict]
+ """The encryption spec to be used for the Agent Engine."""
+
 
 CreateAgentEngineConfigOrDict = Union[
  CreateAgentEngineConfig, CreateAgentEngineConfigDict
@@ -4654,6 +4661,10 @@ class UpdateAgentEngineConfig(_common.BaseModel):
  Recommended value: 2 * cpu + 1. Defaults to 9.
  """,
  )
+ encryption_spec: Optional[EncryptionSpec] = Field(
+ default=None,
+ description="""The encryption spec to be used for the Agent Engine.""",
+ )
  update_mask: Optional[str] = Field(
  default=None,
  description="""The update mask to apply. For the `FieldMask` definition, see
@@ -4710,6 +4721,9 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
  Recommended value: 2 * cpu + 1. Defaults to 9.
  """
 
+ encryption_spec: Optional[EncryptionSpecDict]
+ """The encryption spec to be used for the Agent Engine."""
+
  update_mask: Optional[str]
  """The update mask to apply. For the `FieldMask` definition, see
  https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
@@ -8981,6 +8995,10 @@ class AgentEngineConfig(_common.BaseModel):
  Recommended value: 2 * cpu + 1. Defaults to 9.
  """,
  )
+ encryption_spec: Optional[EncryptionSpec] = Field(
+ default=None,
+ description="""The encryption spec to be used for the Agent Engine.""",
+ )
 
 
 class AgentEngineConfigDict(TypedDict, total=False):
@@ -9053,5 +9071,8 @@ class AgentEngineConfigDict(TypedDict, total=False):
  Recommended value: 2 * cpu + 1. Defaults to 9.
  """
 
+ encryption_spec: Optional[EncryptionSpecDict]
+ """The encryption spec to be used for the Agent Engine."""
+
 
 AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]