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]