Skip to content

feat: Add encryption_spec (CMEK) support to EvaluationRun SDK #7053

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 EvaluationRun SDK #7053

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
65 changes: 63 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 @@ -160,6 +163,9 @@ def _CreateEvaluationRunParameters_to_vertex(
[item for item in getv(from_object, ["analysis_configs"])],
)

if getv(from_object, ["encryption_spec"]) is not None:
setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))

if getv(from_object, ["evaluation_experiment"]) is not None:
setv(
to_object,
Expand All @@ -181,6 +187,9 @@ def _CreateEvaluationSetParameters_to_vertex(
if getv(from_object, ["display_name"]) is not None:
setv(to_object, ["displayName"], getv(from_object, ["display_name"]))

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

Expand Down Expand Up @@ -1422,6 +1433,7 @@ def _create_evaluation_run(
] = None,
config: Optional[types.CreateEvaluationRunConfigOrDict] = None,
analysis_configs: Optional[list[types.AnalysisConfigOrDict]] = None,
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
evaluation_experiment: Optional[str] = None,
) -> types.EvaluationRun:
"""
Expand All @@ -1437,6 +1449,7 @@ def _create_evaluation_run(
inference_configs=inference_configs,
config=config,
analysis_configs=analysis_configs,
encryption_spec=encryption_spec,
evaluation_experiment=evaluation_experiment,
)

Expand Down Expand Up @@ -1505,6 +1518,7 @@ 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:
"""
Expand All @@ -1514,6 +1528,7 @@ def _create_evaluation_set(
parameter_model = types._CreateEvaluationSetParameters(
evaluation_items=evaluation_items,
display_name=display_name,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -3236,6 +3251,7 @@ def create_evaluation_run(
loss_analysis_metrics: Optional[list[Union[str, types.MetricOrDict]]] = None,
loss_analysis_configs: Optional[list[types.LossAnalysisConfigOrDict]] = None,
red_teaming_config: Optional[types.RedTeamingAnalysisConfigOrDict] = None,
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
config: Optional[types.CreateEvaluationRunConfigOrDict] = None,
) -> types.EvaluationRun:
"""Creates an EvaluationRun.
Expand Down Expand Up @@ -3430,6 +3446,7 @@ def create_evaluation_run(
inference_configs=resolved_inference_configs,
analysis_configs=resolved_analysis_configs,
labels=resolved_labels,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -3531,6 +3548,7 @@ 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.
Expand All @@ -3539,6 +3557,8 @@ 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.

Expand All @@ -3548,6 +3568,7 @@ def create_evaluation_set(
return self._create_evaluation_set(
evaluation_items=evaluation_items,
display_name=display_name,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -3695,9 +3716,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 +3751,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 +4013,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 +4024,7 @@ async def _create_evaluation_metric(
display_name=display_name,
description=description,
metric=metric,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -4066,6 +4103,7 @@ async def _create_evaluation_run(
] = None,
config: Optional[types.CreateEvaluationRunConfigOrDict] = None,
analysis_configs: Optional[list[types.AnalysisConfigOrDict]] = None,
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
evaluation_experiment: Optional[str] = None,
) -> types.EvaluationRun:
"""
Expand All @@ -4081,6 +4119,7 @@ async def _create_evaluation_run(
inference_configs=inference_configs,
config=config,
analysis_configs=analysis_configs,
encryption_spec=encryption_spec,
evaluation_experiment=evaluation_experiment,
)

Expand Down Expand Up @@ -4151,6 +4190,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:
"""
Expand All @@ -4160,6 +4200,7 @@ async def _create_evaluation_set(
parameter_model = types._CreateEvaluationSetParameters(
evaluation_items=evaluation_items,
display_name=display_name,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -5515,6 +5556,7 @@ async def create_evaluation_run(
labels: Optional[dict[str, str]] = None,
loss_analysis_metrics: Optional[list[Union[str, types.MetricOrDict]]] = None,
loss_analysis_configs: Optional[list[types.LossAnalysisConfigOrDict]] = None,
encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
config: Optional[types.CreateEvaluationRunConfigOrDict] = None,
) -> types.EvaluationRun:
"""Creates an EvaluationRun.
Expand Down Expand Up @@ -5710,6 +5752,7 @@ async def create_evaluation_run(
inference_configs=resolved_inference_configs,
analysis_configs=resolved_analysis_configs,
labels=resolved_labels,
encryption_spec=encryption_spec,
config=config,
)

Expand Down Expand Up @@ -5816,6 +5859,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.
Expand All @@ -5824,6 +5868,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.

Expand All @@ -5833,6 +5879,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
Expand Down Expand Up @@ -5981,9 +6028,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 +6061,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
42 changes: 42 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,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 @@ -2322,6 +2326,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 @@ -2446,6 +2453,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 @@ -2463,6 +2474,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 Expand Up @@ -3017,6 +3031,10 @@ class _CreateEvaluationRunParameters(_common.BaseModel):
analysis_configs: Optional[list[AnalysisConfig]] = Field(
default=None, description=""""""
)
encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
default=None,
description="""Customer-managed encryption key spec for this EvaluationRun. If set, this EvaluationRun will be secured by this key.""",
)
evaluation_experiment: Optional[str] = Field(
default=None,
description="""The resource name of the parent EvaluationExperiment that this run
Expand Down Expand Up @@ -3052,6 +3070,9 @@ class _CreateEvaluationRunParametersDict(TypedDict, total=False):
analysis_configs: Optional[list[AnalysisConfigDict]]
""""""

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

evaluation_experiment: Optional[str]
"""The resource name of the parent EvaluationExperiment that this run
belongs to. Format:
Expand Down Expand Up @@ -3944,6 +3965,10 @@ class EvaluationRun(_common.BaseModel):
default=None,
description="""The analysis configurations for the evaluation run.""",
)
encryption_spec: Optional[genai_types.EncryptionSpec] = Field(
default=None,
description="""Customer-managed encryption key spec for this EvaluationRun. If set, this EvaluationRun will be secured by this key.""",
)

# TODO(b/448806531): Remove all the overridden _from_response methods once the
# ticket is resolved and published.
Expand Down Expand Up @@ -4047,6 +4072,9 @@ class EvaluationRunDict(TypedDict, total=False):
analysis_configs: Optional[list[AnalysisConfigDict]]
"""The analysis configurations for the evaluation run."""

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


EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict]

Expand Down Expand Up @@ -4076,6 +4104,10 @@ class _CreateEvaluationSetParameters(_common.BaseModel):

evaluation_items: Optional[list[str]] = Field(default=None, description="""""")
display_name: Optional[str] = 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.""",
)
config: Optional[CreateEvaluationSetConfig] = Field(
default=None, description=""""""
)
Expand All @@ -4090,6 +4122,9 @@ class _CreateEvaluationSetParametersDict(TypedDict, total=False):
display_name: Optional[str]
""""""

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

config: Optional[CreateEvaluationSetConfigDict]
""""""

Expand Down Expand Up @@ -4121,6 +4156,10 @@ 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):
Expand All @@ -4144,6 +4183,9 @@ class EvaluationSetDict(TypedDict, total=False):
metadata: Optional[dict[str, Any]]
"""The metadata of the evaluation set."""

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


EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict]

Expand Down
Loading
Loading