Skip to content

Commit c5b69be

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Add encryption_spec (CMEK) support to EvaluationMetric SDK.
PiperOrigin-RevId: 959723760
1 parent 2a18bff commit c5b69be

5 files changed

Lines changed: 112 additions & 2 deletions

File tree

‎agentplatform/_genai/evals.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,9 @@ def _CreateEvaluationMetricParameters_to_vertex(
108108
t.t_metric_for_registry(getv(from_object, ["metric"])),
109109
)
110110

111+
if getv(from_object, ["encryption_spec"]) is not None:
112+
setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
113+
111114
if getv(from_object, ["config"]) is not None:
112115
setv(to_object, ["config"], getv(from_object, ["config"]))
113116

@@ -415,6 +418,9 @@ def _EvaluationMetric_from_vertex(
415418
_UnifiedMetric_from_vertex(getv(from_object, ["metric"]), to_object),
416419
)
417420

421+
if getv(from_object, ["encryptionSpec"]) is not None:
422+
setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"]))
423+
418424
return to_object
419425

420426

@@ -1336,6 +1342,7 @@ def _create_evaluation_metric(
13361342
display_name: Optional[str] = None,
13371343
description: Optional[str] = None,
13381344
metric: Optional[types.MetricOrDict] = None,
1345+
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
13391346
config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
13401347
) -> types.EvaluationMetric:
13411348
"""
@@ -1346,6 +1353,7 @@ def _create_evaluation_metric(
13461353
display_name=display_name,
13471354
description=description,
13481355
metric=metric,
1356+
encryption_spec=encryption_spec,
13491357
config=config,
13501358
)
13511359

@@ -3978,6 +3986,7 @@ async def _create_evaluation_metric(
39783986
display_name: Optional[str] = None,
39793987
description: Optional[str] = None,
39803988
metric: Optional[types.MetricOrDict] = None,
3989+
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
39813990
config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
39823991
) -> types.EvaluationMetric:
39833992
"""
@@ -3988,6 +3997,7 @@ async def _create_evaluation_metric(
39883997
display_name=display_name,
39893998
description=description,
39903999
metric=metric,
4000+
encryption_spec=encryption_spec,
39914001
config=config,
39924002
)
39934003

‎agentplatform/_genai/types/common.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2301,6 +2301,11 @@ class _CreateEvaluationMetricParameters(_common.BaseModel):
23012301
default=None,
23022302
description="""The metric configuration of the evaluation metric.""",
23032303
)
2304+
encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
2305+
default=None,
2306+
description="""Customer-managed encryption key spec for this EvaluationMetric.
2307+
If set, this EvaluationMetric will be secured by this key.""",
2308+
)
23042309
config: Optional[CreateEvaluationMetricConfig] = Field(
23052310
default=None, description=""""""
23062311
)
@@ -2322,6 +2327,10 @@ class _CreateEvaluationMetricParametersDict(TypedDict, total=False):
23222327
metric: Optional[MetricDict]
23232328
"""The metric configuration of the evaluation metric."""
23242329

2330+
encryption_spec: Optional[genai_types.EncryptionSpec]
2331+
"""Customer-managed encryption key spec for this EvaluationMetric.
2332+
If set, this EvaluationMetric will be secured by this key."""
2333+
23252334
config: Optional[CreateEvaluationMetricConfigDict]
23262335
""""""
23272336

@@ -2446,6 +2455,11 @@ class EvaluationMetric(_common.BaseModel):
24462455
default=None,
24472456
description="""The metric configuration of the evaluation metric.""",
24482457
)
2458+
encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
2459+
default=None,
2460+
description="""Customer-managed encryption key spec for this EvaluationMetric.
2461+
If set, this EvaluationMetric will be secured by this key.""",
2462+
)
24492463

24502464

24512465
class EvaluationMetricDict(TypedDict, total=False):
@@ -2463,6 +2477,10 @@ class EvaluationMetricDict(TypedDict, total=False):
24632477
metric: Optional[UnifiedMetricDict]
24642478
"""The metric configuration of the evaluation metric."""
24652479

2480+
encryption_spec: Optional[genai_types.EncryptionSpec]
2481+
"""Customer-managed encryption key spec for this EvaluationMetric.
2482+
If set, this EvaluationMetric will be secured by this key."""
2483+
24662484

24672485
EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict]
24682486

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from tests.unit.agentplatform.genai.replays import pytest_helper
1919
from agentplatform._genai import types
2020
from google.genai import errors
21+
from google.genai import types as genai_types
2122
import pytest
2223

2324

@@ -69,6 +70,35 @@ def test_list_evaluation_metrics_with_filter(client):
6970
assert metric.display_name == "tone-check-v1"
7071

7172

73+
def test_create_evaluation_metric_with_cmek(client):
74+
"""CMEK: encryption_spec is forwarded in the request and returned on GET."""
75+
client._api_client._http_options.api_version = "v1beta1"
76+
_KMS_KEY = (
77+
"projects/977012026409/locations/us-central1"
78+
"/keyRings/test-kr/cryptoKeys/test-key"
79+
)
80+
result = client.evals.create_evaluation_metric(
81+
display_name="test_cmek_metric",
82+
description="test_cmek_description",
83+
metric=types.LLMMetric(
84+
name="custom_llm_metric", prompt_template="test_prompt_template"
85+
),
86+
encryption_spec=genai_types.EncryptionSpec(kms_key_name=_KMS_KEY),
87+
)
88+
assert isinstance(result, str)
89+
assert re.match(
90+
r"^projects/[^/]+/locations/[^/]+/evaluationMetrics/[^/]+$",
91+
result,
92+
)
93+
94+
metric = client.evals.get_evaluation_metric(metric_resource_name=result)
95+
assert isinstance(metric, types.EvaluationMetric)
96+
assert metric.display_name == "test_cmek_metric"
97+
# encryption_spec is returned from the API and surfaced on the resource.
98+
assert metric.encryption_spec is not None
99+
assert metric.encryption_spec.kms_key_name == _KMS_KEY
100+
101+
72102
# The setup function registers the module and method for the recorder
73103
pytestmark = pytest_helper.setup(
74104
file=__file__,

‎vertexai/_genai/evals.py‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,9 @@ def _CreateEvaluationMetricParameters_to_vertex(
8585
t.t_metric_for_registry(getv(from_object, ["metric"])),
8686
)
8787

88+
if getv(from_object, ["encryption_spec"]) is not None:
89+
setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
90+
8891
if getv(from_object, ["config"]) is not None:
8992
setv(to_object, ["config"], getv(from_object, ["config"]))
9093

@@ -408,6 +411,9 @@ def _EvaluationMetric_from_vertex(
408411
_UnifiedMetric_from_vertex(getv(from_object, ["metric"]), to_object),
409412
)
410413

414+
if getv(from_object, ["encryptionSpec"]) is not None:
415+
setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"]))
416+
411417
return to_object
412418

413419

@@ -1226,6 +1232,7 @@ def _create_evaluation_metric(
12261232
display_name: Optional[str] = None,
12271233
description: Optional[str] = None,
12281234
metric: Optional[types.MetricOrDict] = None,
1235+
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
12291236
config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
12301237
) -> types.EvaluationMetric:
12311238
"""
@@ -1236,6 +1243,7 @@ def _create_evaluation_metric(
12361243
display_name=display_name,
12371244
description=description,
12381245
metric=metric,
1246+
encryption_spec=encryption_spec,
12391247
config=config,
12401248
)
12411249

@@ -3160,9 +3168,22 @@ def create_evaluation_metric(
31603168
display_name: Optional[str] = None,
31613169
description: Optional[str] = None,
31623170
metric: Optional[types.MetricOrDict] = None,
3171+
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
31633172
config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
31643173
) -> str:
3165-
"""Creates an EvaluationMetric."""
3174+
"""Creates an EvaluationMetric.
3175+
3176+
Args:
3177+
display_name: The display name of the EvaluationMetric.
3178+
description: The description of the EvaluationMetric.
3179+
metric: The metric configuration.
3180+
encryption_spec: Customer-managed encryption key spec. If set, this
3181+
EvaluationMetric will be secured by the provided key.
3182+
config: Optional configuration for the request.
3183+
3184+
Returns:
3185+
The resource name of the created EvaluationMetric.
3186+
"""
31663187
if metric and not isinstance(metric, dict):
31673188
# metric is now Metric | LazyLoadedPrebuiltMetric (RubricMetric)
31683189
# Mypy correctly narrows the type here, so cast is not needed.
@@ -3182,6 +3203,7 @@ def create_evaluation_metric(
31823203
display_name=display_name,
31833204
description=description,
31843205
metric=metric,
3206+
encryption_spec=encryption_spec,
31853207
config=config,
31863208
)
31873209
# result.name is Optional[str], but we know it's always returned on creation
@@ -3776,6 +3798,7 @@ async def _create_evaluation_metric(
37763798
display_name: Optional[str] = None,
37773799
description: Optional[str] = None,
37783800
metric: Optional[types.MetricOrDict] = None,
3801+
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
37793802
config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
37803803
) -> types.EvaluationMetric:
37813804
"""
@@ -3786,6 +3809,7 @@ async def _create_evaluation_metric(
37863809
display_name=display_name,
37873810
description=description,
37883811
metric=metric,
3812+
encryption_spec=encryption_spec,
37893813
config=config,
37903814
)
37913815

@@ -5361,9 +5385,22 @@ async def create_evaluation_metric(
53615385
display_name: Optional[str] = None,
53625386
description: Optional[str] = None,
53635387
metric: Optional[types.MetricOrDict] = None,
5388+
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
53645389
config: Optional[types.CreateEvaluationMetricConfigOrDict] = None,
53655390
) -> str:
5366-
"""Creates an EvaluationMetric."""
5391+
"""Creates an EvaluationMetric.
5392+
5393+
Args:
5394+
display_name: The display name of the EvaluationMetric.
5395+
description: The description of the EvaluationMetric.
5396+
metric: The metric configuration.
5397+
encryption_spec: Customer-managed encryption key spec. If set, this
5398+
EvaluationMetric will be secured by the provided key.
5399+
config: Optional configuration for the request.
5400+
5401+
Returns:
5402+
The resource name of the created EvaluationMetric.
5403+
"""
53675404
if metric and not isinstance(metric, dict):
53685405
resolved_metrics = _evals_common._resolve_metrics(
53695406
[metric], self._api_client
@@ -5381,6 +5418,7 @@ async def create_evaluation_metric(
53815418
display_name=display_name,
53825419
description=description,
53835420
metric=metric,
5421+
encryption_spec=encryption_spec,
53845422
config=config,
53855423
)
53865424
return cast(str, result.name)

‎vertexai/_genai/types/common.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2115,6 +2115,10 @@ class _CreateEvaluationMetricParameters(_common.BaseModel):
21152115
default=None,
21162116
description="""The metric configuration of the evaluation metric.""",
21172117
)
2118+
encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
2119+
default=None,
2120+
description="""Customer-managed encryption key spec for this EvaluationMetric. If set, this EvaluationMetric will be secured by this key.""",
2121+
)
21182122
config: Optional[CreateEvaluationMetricConfig] = Field(
21192123
default=None, description=""""""
21202124
)
@@ -2136,6 +2140,9 @@ class _CreateEvaluationMetricParametersDict(TypedDict, total=False):
21362140
metric: Optional[MetricDict]
21372141
"""The metric configuration of the evaluation metric."""
21382142

2143+
encryption_spec: Optional[genai_types.EncryptionSpecDict]
2144+
"""Customer-managed encryption key spec for this EvaluationMetric. If set, this EvaluationMetric will be secured by this key."""
2145+
21392146
config: Optional[CreateEvaluationMetricConfigDict]
21402147
""""""
21412148

@@ -2260,6 +2267,10 @@ class EvaluationMetric(_common.BaseModel):
22602267
default=None,
22612268
description="""The metric configuration of the evaluation metric.""",
22622269
)
2270+
encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
2271+
default=None,
2272+
description="""Customer-managed encryption key spec for this EvaluationMetric. If set, this EvaluationMetric will be secured by this key.""",
2273+
)
22632274

22642275

22652276
class EvaluationMetricDict(TypedDict, total=False):
@@ -2277,6 +2288,9 @@ class EvaluationMetricDict(TypedDict, total=False):
22772288
metric: Optional[UnifiedMetricDict]
22782289
"""The metric configuration of the evaluation metric."""
22792290

2291+
encryption_spec: Optional[genai_types.EncryptionSpecDict]
2292+
"""Customer-managed encryption key spec for this EvaluationMetric. If set, this EvaluationMetric will be secured by this key."""
2293+
22802294

22812295
EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict]
22822296

0 commit comments

Comments
 (0)