From 11d0c0d4924d90e2bd9e8997cd430ddc30a43191 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Wed, 5 Aug 2026 10:04:32 -0700 Subject: [PATCH] feat: Add encryption_spec (CMEK) support to EvaluationMetric SDK. PiperOrigin-RevId: 959723760 --- agentplatform/_genai/evals.py | 14 +++++++++ agentplatform/_genai/types/common.py | 18 +++++++++++ .../genai/replays/test_evaluation_metric.py | 30 +++++++++++++++++++ 3 files changed, 62 insertions(+) diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index 4b600f7054..d408ced686 100644 --- a/agentplatform/_genai/evals.py +++ b/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"])) @@ -415,6 +418,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 +1342,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 +1353,7 @@ def _create_evaluation_metric( display_name=display_name, description=description, metric=metric, + encryption_spec=encryption_spec, config=config, ) @@ -3695,6 +3703,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 +3726,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 +3988,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 +3999,7 @@ async def _create_evaluation_metric( display_name=display_name, description=description, metric=metric, + encryption_spec=encryption_spec, config=config, ) @@ -5981,6 +5993,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 +6014,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) diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index b0343429c1..549dc9f85e 100644 --- a/agentplatform/_genai/types/common.py +++ b/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] diff --git a/tests/unit/agentplatform/genai/replays/test_evaluation_metric.py b/tests/unit/agentplatform/genai/replays/test_evaluation_metric.py index c9443f9395..40a64ba209 100644 --- a/tests/unit/agentplatform/genai/replays/test_evaluation_metric.py +++ b/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__,