Skip to content

feat: Add encryption_spec (CMEK) support to EvaluationMetric SDK. #7051

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Add encryption_spec (CMEK) support to EvaluationMetric SDK. #7051

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 40 additions & 2 deletions agentplatform/_genai/evals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"]))

Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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:
"""
Expand All @@ -1346,6 +1353,7 @@ def _create_evaluation_metric(
display_name=display_name,
description=description,
metric=metric,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -3695,9 +3703,22 @@ 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."""
"""Creates an EvaluationMetric.

Args:
display_name: The display name of the EvaluationMetric.
description: The description of the EvaluationMetric.
metric: The metric configuration.
encryption_spec: Customer-managed encryption key spec. If set, this
EvaluationMetric will be secured by the provided key.
config: Optional configuration for the request.

Returns:
The resource name of the created EvaluationMetric.
"""
if metric and not isinstance(metric, dict):
# metric is now Metric | LazyLoadedPrebuiltMetric (RubricMetric)
# Mypy correctly narrows the type here, so cast is not needed.
Expand All @@ -3717,6 +3738,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
Expand Down Expand Up @@ -3978,6 +4000,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:
"""
Expand All @@ -3988,6 +4011,7 @@ async def _create_evaluation_metric(
display_name=display_name,
description=description,
metric=metric,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -5981,9 +6005,22 @@ 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."""
"""Creates an EvaluationMetric.

Args:
display_name: The display name of the EvaluationMetric.
description: The description of the EvaluationMetric.
metric: The metric configuration.
encryption_spec: Customer-managed encryption key spec. If set, this
EvaluationMetric will be secured by the provided key.
config: Optional configuration for the request.

Returns:
The resource name of the created EvaluationMetric.
"""
if metric and not isinstance(metric, dict):
resolved_metrics = _evals_common._resolve_metrics(
[metric], self._api_client
Expand All @@ -6001,6 +6038,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)
Expand Down
18 changes: 18 additions & 0 deletions agentplatform/_genai/types/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=""""""
)
Expand All @@ -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]
""""""

Expand Down Expand Up @@ -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):
Expand All @@ -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]

Expand Down
30 changes: 30 additions & 0 deletions tests/unit/agentplatform/genai/replays/test_evaluation_metric.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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__,
Expand Down
42 changes: 40 additions & 2 deletions vertexai/_genai/evals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,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"]))

Expand Down Expand Up @@ -408,6 +411,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


Expand Down Expand Up @@ -1226,6 +1232,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:
"""
Expand All @@ -1236,6 +1243,7 @@ def _create_evaluation_metric(
display_name=display_name,
description=description,
metric=metric,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -3160,9 +3168,22 @@ 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."""
"""Creates an EvaluationMetric.

Args:
display_name: The display name of the EvaluationMetric.
description: The description of the EvaluationMetric.
metric: The metric configuration.
encryption_spec: Customer-managed encryption key spec. If set, this
EvaluationMetric will be secured by the provided key.
config: Optional configuration for the request.

Returns:
The resource name of the created EvaluationMetric.
"""
if metric and not isinstance(metric, dict):
# metric is now Metric | LazyLoadedPrebuiltMetric (RubricMetric)
# Mypy correctly narrows the type here, so cast is not needed.
Expand All @@ -3182,6 +3203,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
Expand Down Expand Up @@ -3776,6 +3798,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:
"""
Expand All @@ -3786,6 +3809,7 @@ async def _create_evaluation_metric(
display_name=display_name,
description=description,
metric=metric,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -5361,9 +5385,22 @@ 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."""
"""Creates an EvaluationMetric.

Args:
display_name: The display name of the EvaluationMetric.
description: The description of the EvaluationMetric.
metric: The metric configuration.
encryption_spec: Customer-managed encryption key spec. If set, this
EvaluationMetric will be secured by the provided key.
config: Optional configuration for the request.

Returns:
The resource name of the created EvaluationMetric.
"""
if metric and not isinstance(metric, dict):
resolved_metrics = _evals_common._resolve_metrics(
[metric], self._api_client
Expand All @@ -5381,6 +5418,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)
Expand Down
14 changes: 14 additions & 0 deletions vertexai/_genai/types/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2115,6 +2115,10 @@ 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=""""""
)
Expand All @@ -2136,6 +2140,9 @@ class _CreateEvaluationMetricParametersDict(TypedDict, total=False):
metric: Optional[MetricDict]
"""The metric configuration of the evaluation metric."""

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

config: Optional[CreateEvaluationMetricConfigDict]
""""""

Expand Down Expand Up @@ -2260,6 +2267,10 @@ 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):
Expand All @@ -2277,6 +2288,9 @@ class EvaluationMetricDict(TypedDict, total=False):
metric: Optional[UnifiedMetricDict]
"""The metric configuration of the evaluation metric."""

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


EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict]

Expand Down