feat: Add encryption_spec (CMEK) support to EvaluationSet resource

agentplatform/_genai/evals.py

@@ -108,6 +108,9 @@ def _CreateEvaluationMetricParameters_to_vertex(
             t.t_metric_for_registry(getv(from_object, ["metric"])),
         )
 
+    if getv(from_object, ["encryption_spec"]) is not None:
+        setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
+
     if getv(from_object, ["config"]) is not None:
         setv(to_object, ["config"], getv(from_object, ["config"]))
 
@@ -184,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
 
 
@@ -415,6 +421,9 @@ def _EvaluationMetric_from_vertex(
             _UnifiedMetric_from_vertex(getv(from_object, ["metric"]), to_object),
         )
 
+    if getv(from_object, ["encryptionSpec"]) is not None:
+        setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"]))
+
     return to_object
 
 
@@ -1336,6 +1345,7 @@ def _create_evaluation_metric(
         display_name: Optional[str] = None,
         description: Optional[str] = None,
         metric: Optional[types.MetricOrDict] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
         config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
     ) -> types.EvaluationMetric:
         """
@@ -1346,6 +1356,7 @@ def _create_evaluation_metric(
             display_name=display_name,
             description=description,
             metric=metric,
+            encryption_spec=encryption_spec,
             config=config,
         )
 
@@ -1506,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.
@@ -1515,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]]
@@ -3519,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.
@@ -3539,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,
@@ -3695,6 +3714,7 @@ def create_evaluation_metric(
         display_name: Optional[str] = None,
         description: Optional[str] = None,
         metric: Optional[types.MetricOrDict] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
         config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
     ) -> str:
         """Creates an EvaluationMetric."""
@@ -3717,6 +3737,7 @@ def create_evaluation_metric(
             display_name=display_name,
             description=description,
             metric=metric,
+            encryption_spec=encryption_spec,
             config=config,
         )
         # result.name is Optional[str], but we know it's always returned on creation
@@ -3978,6 +3999,7 @@ async def _create_evaluation_metric(
         display_name: Optional[str] = None,
         description: Optional[str] = None,
         metric: Optional[types.MetricOrDict] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
         config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
     ) -> types.EvaluationMetric:
         """
@@ -3988,6 +4010,7 @@ async def _create_evaluation_metric(
             display_name=display_name,
             description=description,
             metric=metric,
+            encryption_spec=encryption_spec,
             config=config,
         )
 
@@ -4152,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.
@@ -4161,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]]
@@ -5816,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.
@@ -5824,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.
 
@@ -5833,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
@@ -5981,6 +6010,7 @@ async def create_evaluation_metric(
         display_name: Optional[str] = None,
         description: Optional[str] = None,
         metric: Optional[types.MetricOrDict] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
         config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
     ) -> str:
         """Creates an EvaluationMetric."""
@@ -6001,6 +6031,7 @@ async def create_evaluation_metric(
             display_name=display_name,
             description=description,
             metric=metric,
+            encryption_spec=encryption_spec,
             config=config,
         )
         return cast(str, result.name)

agentplatform/_genai/types/common.py

@@ -2301,6 +2301,11 @@ class _CreateEvaluationMetricParameters(_common.BaseModel):
         default=None,
         description="""The metric configuration of the evaluation metric.""",
     )
+    encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
+        default=None,
+        description="""Customer-managed encryption key spec for this EvaluationMetric.
+      If set, this EvaluationMetric will be secured by this key.""",
+    )
     config: Optional[CreateEvaluationMetricConfig] = Field(
         default=None, description=""""""
     )
@@ -2322,6 +2327,10 @@ class _CreateEvaluationMetricParametersDict(TypedDict, total=False):
     metric: Optional[MetricDict]
     """The metric configuration of the evaluation metric."""
 
+    encryption_spec: Optional[genai_types.EncryptionSpec]
+    """Customer-managed encryption key spec for this EvaluationMetric.
+      If set, this EvaluationMetric will be secured by this key."""
+
     config: Optional[CreateEvaluationMetricConfigDict]
     """"""
 
@@ -2446,6 +2455,11 @@ class EvaluationMetric(_common.BaseModel):
         default=None,
         description="""The metric configuration of the evaluation metric.""",
     )
+    encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
+        default=None,
+        description="""Customer-managed encryption key spec for this EvaluationMetric.
+      If set, this EvaluationMetric will be secured by this key.""",
+    )
 
 
 class EvaluationMetricDict(TypedDict, total=False):
@@ -2463,6 +2477,10 @@ class EvaluationMetricDict(TypedDict, total=False):
     metric: Optional[UnifiedMetricDict]
     """The metric configuration of the evaluation metric."""
 
+    encryption_spec: Optional[genai_types.EncryptionSpec]
+    """Customer-managed encryption key spec for this EvaluationMetric.
+      If set, this EvaluationMetric will be secured by this key."""
+
 
 EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict]
 
@@ -4079,6 +4097,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):
@@ -4093,6 +4116,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
@@ -4121,6 +4148,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):
@@ -4144,6 +4176,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",)
 
 

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

@@ -18,6 +18,7 @@
 from tests.unit.agentplatform.genai.replays import pytest_helper
 from agentplatform._genai import types
 from google.genai import errors
+from google.genai import types as genai_types
 import pytest
 
 
@@ -69,6 +70,35 @@ def test_list_evaluation_metrics_with_filter(client):
         assert metric.display_name == "tone-check-v1"
 
 
+def test_create_evaluation_metric_with_cmek(client):
+    """CMEK: encryption_spec is forwarded in the request and returned on GET."""
+    client._api_client._http_options.api_version = "v1beta1"
+    _KMS_KEY = (
+        "projects/977012026409/locations/us-central1"
+        "/keyRings/test-kr/cryptoKeys/test-key"
+    )
+    result = client.evals.create_evaluation_metric(
+        display_name="test_cmek_metric",
+        description="test_cmek_description",
+        metric=types.LLMMetric(
+            name="custom_llm_metric", prompt_template="test_prompt_template"
+        ),
+        encryption_spec=genai_types.EncryptionSpec(kms_key_name=_KMS_KEY),
+    )
+    assert isinstance(result, str)
+    assert re.match(
+        r"^projects/[^/]+/locations/[^/]+/evaluationMetrics/[^/]+$",
+        result,
+    )
+
+    metric = client.evals.get_evaluation_metric(metric_resource_name=result)
+    assert isinstance(metric, types.EvaluationMetric)
+    assert metric.display_name == "test_cmek_metric"
+    # encryption_spec is returned from the API and surfaced on the resource.
+    assert metric.encryption_spec is not None
+    assert metric.encryption_spec.kms_key_name == _KMS_KEY
+
+
 # The setup function registers the module and method for the recorder
 pytestmark = pytest_helper.setup(
     file=__file__,