feat: Add encryption_spec (CMEK) support to EvaluationSet resource

agentplatform/_genai/evals.py

@@ -187,6 +187,9 @@ def _CreateEvaluationSetParameters_to_vertex(
     if getv(from_object, ["config"]) is not None:
         setv(to_object, ["config"], getv(from_object, ["config"]))
 
+    if getv(from_object, ["encryption_spec"]) is not None:
+        setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
+
     return to_object
 
 
@@ -1514,6 +1517,7 @@ def _create_evaluation_set(
         evaluation_items: list[str],
         display_name: Optional[str] = None,
         config: Optional[types.CreateEvaluationSetConfigOrDict] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
     ) -> types.EvaluationSet:
         """
         Creates an EvaluationSet.
@@ -1523,6 +1527,7 @@ def _create_evaluation_set(
             evaluation_items=evaluation_items,
             display_name=display_name,
             config=config,
+            encryption_spec=encryption_spec,
         )
 
         request_url_dict: Optional[dict[str, str]]
@@ -3527,18 +3532,20 @@ def create_evaluation_item(
         Returns:
           The evaluation item.
         """
-        return self._create_evaluation_item(
+        result = self._create_evaluation_item(
             evaluation_item_type=evaluation_item_type,
             gcs_uri=gcs_uri,
             display_name=display_name,
             config=config,
         )
+        return result
 
     def create_evaluation_set(
         self,
         *,
         evaluation_items: list[str],
         display_name: Optional[str] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
         config: Optional[types.CreateEvaluationSetConfigOrDict] = None,
     ) -> types.EvaluationSet:
         """Creates an EvaluationSet.
@@ -3547,17 +3554,21 @@ def create_evaluation_set(
           evaluation_items: The list of evaluation item names. Format:
             `projects/{project}/locations/{location}/evaluationItems/{evaluation_item}`
           display_name: The display name of the evaluation set.
+          encryption_spec: Customer-managed encryption key spec. If set, this
+            EvaluationSet will be secured by the provided key.
           config: The optional configuration for the evaluation set. Must be a dict or
               `types.CreateEvaluationSetConfigOrDict` type.
 
         Returns:
           The evaluation set.
         """
-        return self._create_evaluation_set(
+        result = self._create_evaluation_set(
             evaluation_items=evaluation_items,
             display_name=display_name,
+            encryption_spec=encryption_spec,
             config=config,
         )
+        return result
 
     def generate_conversation_scenarios(
         self,
@@ -4164,6 +4175,7 @@ async def _create_evaluation_set(
         evaluation_items: list[str],
         display_name: Optional[str] = None,
         config: Optional[types.CreateEvaluationSetConfigOrDict] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
     ) -> types.EvaluationSet:
         """
         Creates an EvaluationSet.
@@ -4173,6 +4185,7 @@ async def _create_evaluation_set(
             evaluation_items=evaluation_items,
             display_name=display_name,
             config=config,
+            encryption_spec=encryption_spec,
         )
 
         request_url_dict: Optional[dict[str, str]]
@@ -5828,6 +5841,7 @@ async def create_evaluation_set(
         *,
         evaluation_items: list[str],
         display_name: Optional[str] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
         config: Optional[types.CreateEvaluationSetConfigOrDict] = None,
     ) -> types.EvaluationSet:
         """Creates an EvaluationSet.
@@ -5836,6 +5850,8 @@ async def create_evaluation_set(
           evaluation_items: The list of evaluation item names. Format:
             `projects/{project}/locations/{location}/evaluationItems/{evaluation_item}`
           display_name: The display name of the evaluation set.
+          encryption_spec: Customer-managed encryption key spec. If set, this
+            EvaluationSet will be secured by the provided key.
           config: The optional configuration for the evaluation set. Must be a dict or
               `types.CreateEvaluationSetConfigOrDict` type.
 
@@ -5845,6 +5861,7 @@ async def create_evaluation_set(
         result = await self._create_evaluation_set(
             evaluation_items=evaluation_items,
             display_name=display_name,
+            encryption_spec=encryption_spec,
             config=config,
         )
         return result

agentplatform/_genai/types/common.py

@@ -4250,6 +4250,11 @@ class _CreateEvaluationSetParameters(_common.BaseModel):
     config: Optional[CreateEvaluationSetConfig] = Field(
         default=None, description=""""""
     )
+    encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
+        default=None,
+        description="""Customer-managed encryption key spec for this EvaluationSet.
+      If set, this EvaluationSet will be secured by this key.""",
+    )
 
 
 class _CreateEvaluationSetParametersDict(TypedDict, total=False):
@@ -4264,6 +4269,10 @@ class _CreateEvaluationSetParametersDict(TypedDict, total=False):
     config: Optional[CreateEvaluationSetConfigDict]
     """"""
 
+    encryption_spec: Optional[genai_types.EncryptionSpec]
+    """Customer-managed encryption key spec for this EvaluationSet.
+      If set, this EvaluationSet will be secured by this key."""
+
 
 _CreateEvaluationSetParametersOrDict = Union[
     _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict
@@ -4292,6 +4301,11 @@ class EvaluationSet(_common.BaseModel):
     metadata: Optional[dict[str, Any]] = Field(
         default=None, description="""The metadata of the evaluation set."""
     )
+    encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
+        default=None,
+        description="""Customer-managed encryption key spec for this EvaluationSet.
+      If set, this EvaluationSet will be secured by this key.""",
+    )
 
 
 class EvaluationSetDict(TypedDict, total=False):
@@ -4315,6 +4329,10 @@ class EvaluationSetDict(TypedDict, total=False):
     metadata: Optional[dict[str, Any]]
     """The metadata of the evaluation set."""
 
+    encryption_spec: Optional[genai_types.EncryptionSpec]
+    """Customer-managed encryption key spec for this EvaluationSet.
+      If set, this EvaluationSet will be secured by this key."""
+
 
 EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict]
 

tests/unit/agentplatform/genai/replays/test_create_evaluation_set.py

@@ -16,6 +16,7 @@
 
 from tests.unit.agentplatform.genai.replays import pytest_helper
 from agentplatform import types
+from google.genai import types as genai_types
 import pytest
 
 
@@ -36,6 +37,26 @@ def test_create_eval_set(client):
     assert evaluation_set.evaluation_items == EVAL_ITEMS
 
 
+_KMS_KEY = (
+    "projects/503583131166/locations/us-central1"
+    "/keyRings/test-kr/cryptoKeys/test-key"
+)
+
+
+def test_create_eval_set_with_cmek(client):
+    """CMEK: encryption_spec is forwarded in the request and returned on the resource."""
+    evaluation_set = client.evals.create_evaluation_set(
+        evaluation_items=EVAL_ITEMS,
+        display_name=DISPLAY_NAME,
+        encryption_spec=genai_types.EncryptionSpec(kms_key_name=_KMS_KEY),
+    )
+    assert isinstance(evaluation_set, types.EvaluationSet)
+    assert evaluation_set.display_name == DISPLAY_NAME
+    assert evaluation_set.evaluation_items == EVAL_ITEMS
+    assert evaluation_set.encryption_spec is not None
+    assert evaluation_set.encryption_spec.kms_key_name == _KMS_KEY
+
+
 pytest_plugins = ("pytest_asyncio",)