diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index 4b600f7054..0a90627bc3 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"])) @@ -167,6 +170,9 @@ def _CreateEvaluationRunParameters_to_vertex( getv(from_object, ["evaluation_experiment"]), ) + if getv(from_object, ["encryption_spec"]) is not None: + setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + return to_object @@ -184,6 +190,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 +424,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 @@ -668,6 +680,9 @@ def _EvaluationRun_from_vertex( [item for item in getv(from_object, ["analysisConfigs"])], ) + if getv(from_object, ["encryptionSpec"]) is not None: + setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) + return to_object @@ -1336,6 +1351,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 +1362,7 @@ def _create_evaluation_metric( display_name=display_name, description=description, metric=metric, + encryption_spec=encryption_spec, config=config, ) @@ -1423,6 +1440,7 @@ def _create_evaluation_run( config: Optional[types.CreateEvaluationRunConfigOrDict] = None, analysis_configs: Optional[list[types.AnalysisConfigOrDict]] = None, evaluation_experiment: Optional[str] = None, + encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None, ) -> types.EvaluationRun: """ Creates an EvaluationRun. @@ -1438,6 +1456,7 @@ def _create_evaluation_run( config=config, analysis_configs=analysis_configs, evaluation_experiment=evaluation_experiment, + encryption_spec=encryption_spec, ) request_url_dict: Optional[dict[str, str]] @@ -1506,6 +1525,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 +1535,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]] @@ -3236,69 +3257,70 @@ 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. Args: - dataset: The dataset to evaluate. Either an EvaluationRunDataSource or an EvaluationDataset. - dest: The GCS URI prefix to write the evaluation results to. - metrics: The list of metrics to evaluate. - name: The name of the evaluation run. - display_name: The display name of the evaluation run. - evaluation_experiment: The resource name of an existing - EvaluationExperiment to group this run under. If omitted, a new - EvaluationExperiment is created automatically so the run is visible in - the Agent Platform UI. Pass an existing experiment name to group - multiple runs together. - agent_info: The agent info to evaluate. Mutually exclusive with - `inference_configs`. - agent: The agent resource name in str type. Accepts either an Agent - Engine resource name - `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine_id}` - or a Gemini Agent (Vertex AI Agent) resource name - `projects/{project}/locations/{location}/agents/{agent}`. When a Gemini - Agent resource is provided, the backend scrapes the agent to produce - agent responses. If an Agent Engine resource name is provided, runs - inference with the deployed agent to get agent responses for evaluation. - The `agent` parameter is required if `agent_info` is provided. - user_simulator_config: The user simulator configuration for agent evaluation. - If `agent_info` is provided without `inference_configs`, this config is used - to automatically construct the inference configuration. If not specified, - or if `max_turn` is not set, `max_turn` defaults to 5. - The `model_name` inside this config can be either a full model path or a - short model name, e.g. `gemini-3-preview-flash`. - inference_configs: The candidate to inference config map for the evaluation run. - The key is the candidate name, and the value is the inference config. - If provided, `agent_info` must be None. If omitted and `agent_info` is provided, - this will be automatically constructed using `agent_info` and `user_simulator_config`. - The `model` field of an inference config accepts a short Gemini model - name (e.g. `gemini-2.5-flash`), which is automatically expanded to a - fully-qualified resource name using the client's project and location, - or an already fully-qualified publisher-model or endpoint resource - name. - Example: - {"candidate-1": types.EvaluationRunInferenceConfig(model="gemini-2.5-flash")} - labels: The labels to apply to the evaluation run. - loss_analysis_metrics: This field is experimental and may change in future - versions. Optional list of metrics to run loss analysis on. The - candidate is auto-inferred from ``inference_configs`` or - ``agent_info`` when there is exactly one candidate. Each metric can be - a string (e.g., ``"multi_turn_task_success_v1"``), a ``Metric`` - object, or a ``RubricMetric`` enum - (e.g., ``types.RubricMetric.MULTI_TURN_TASK_SUCCESS``). Loss analysis - runs after metric calculation completes. - Mutually exclusive with ``loss_analysis_configs``. - Example:: - - loss_analysis_metrics=[ - types.RubricMetric.MULTI_TURN_TASK_SUCCESS, - types.RubricMetric.MULTI_TURN_TOOL_USE_QUALITY, - ] - loss_analysis_configs: This field is experimental and may change in future - versions. Optional list of ``LossAnalysisConfig`` objects for full - control over loss analysis, including explicit candidate and - advanced options like ``predefined_taxonomy`` and + dataset: The dataset to evaluate. Either an EvaluationRunDataSource or an EvaluationDataset. + dest: The GCS URI prefix to write the evaluation results to. + metrics: The list of metrics to evaluate. + name: The name of the evaluation run. + display_name: The display name of the evaluation run. + evaluation_experiment: The resource name of an existing + EvaluationExperiment to group this run under. If omitted, a new + EvaluationExperiment is created automatically so the run is visible in + the Agent Platform UI. Pass an existing experiment name to group + multiple runs together. + agent_info: The agent info to evaluate. Mutually exclusive with + `inference_configs`. + agent: The agent resource name in str type. Accepts either an Agent + Engine resource name + `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine_id}` + or a Gemini Agent (Vertex AI Agent) resource name + `projects/{project}/locations/{location}/agents/{agent}`. When a Gemini + Agent resource is provided, the backend scrapes the agent to produce + agent responses. If an Agent Engine resource name is provided, runs + inference with the deployed agent to get agent responses for evaluation. + The `agent` parameter is required if `agent_info` is provided. + user_simulator_config: The user simulator configuration for agent evaluation. + If `agent_info` is provided without `inference_configs`, this config is used + to automatically construct the inference configuration. If not specified, + or if `max_turn` is not set, `max_turn` defaults to 5. + The `model_name` inside this config can be either a full model path or a + short model name, e.g. `gemini-3-preview-flash`. + inference_configs: The candidate to inference config map for the evaluation run. + The key is the candidate name, and the value is the inference config. + If provided, `agent_info` must be None. If omitted and `agent_info` is provided, + this will be automatically constructed using `agent_info` and `user_simulator_config`. + The `model` field of an inference config accepts a short Gemini model + name (e.g. `gemini-2.5-flash`), which is automatically expanded to a + fully-qualified resource name using the client's project and location, + or an already fully-qualified publisher-model or endpoint resource + name. + Example: + {"candidate-1": types.EvaluationRunInferenceConfig(model="gemini-2.5-flash")} + labels: The labels to apply to the evaluation run. + loss_analysis_metrics: This field is experimental and may change in future + versions. Optional list of metrics to run loss analysis on. The + candidate is auto-inferred from ``inference_configs`` or + ``agent_info`` when there is exactly one candidate. Each metric can be + a string (e.g., ``"multi_turn_task_success_v1"``), a ``Metric`` + object, or a ``RubricMetric`` enum + (e.g., ``types.RubricMetric.MULTI_TURN_TASK_SUCCESS``). Loss analysis + runs after metric calculation completes. + Mutually exclusive with ``loss_analysis_configs``. + Example:: + + loss_analysis_metrics=[ + types.RubricMetric.MULTI_TURN_TASK_SUCCESS, + types.RubricMetric.MULTI_TURN_TOOL_USE_QUALITY, + ] + loss_analysis_configs: This field is experimental and may change in future + versions. Optional list of ``LossAnalysisConfig`` objects for full + control over loss analysis, including explicit candidate and + advanced options like ``predefined_taxonomy`` and ``max_top_cluster_count``. Mutually exclusive with ``loss_analysis_metrics``. config: The configuration for the evaluation run. @@ -3430,6 +3452,7 @@ def create_evaluation_run( inference_configs=resolved_inference_configs, analysis_configs=resolved_analysis_configs, labels=resolved_labels, + encryption_spec=encryption_spec, config=config, ) @@ -3519,18 +3542,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 +3564,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 +3724,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 +3747,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 +4009,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 +4020,7 @@ async def _create_evaluation_metric( display_name=display_name, description=description, metric=metric, + encryption_spec=encryption_spec, config=config, ) @@ -4067,6 +4100,7 @@ async def _create_evaluation_run( config: Optional[types.CreateEvaluationRunConfigOrDict] = None, analysis_configs: Optional[list[types.AnalysisConfigOrDict]] = None, evaluation_experiment: Optional[str] = None, + encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None, ) -> types.EvaluationRun: """ Creates an EvaluationRun. @@ -4082,6 +4116,7 @@ async def _create_evaluation_run( config=config, analysis_configs=analysis_configs, evaluation_experiment=evaluation_experiment, + encryption_spec=encryption_spec, ) request_url_dict: Optional[dict[str, str]] @@ -4152,6 +4187,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 +4197,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]] @@ -5515,6 +5552,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. @@ -5710,6 +5748,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, ) @@ -5816,6 +5855,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 +5864,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 +5875,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 +6024,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 +6045,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..1b38648211 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] @@ -3023,6 +3041,11 @@ class _CreateEvaluationRunParameters(_common.BaseModel): belongs to. Format: `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""", ) + 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.""", + ) class _CreateEvaluationRunParametersDict(TypedDict, total=False): @@ -3057,6 +3080,10 @@ class _CreateEvaluationRunParametersDict(TypedDict, total=False): belongs to. Format: `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationRun. + If set, this EvaluationRun will be secured by this key.""" + _CreateEvaluationRunParametersOrDict = Union[ _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict @@ -3944,6 +3971,11 @@ 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. @@ -4047,6 +4079,10 @@ class EvaluationRunDict(TypedDict, total=False): analysis_configs: Optional[list[AnalysisConfigDict]] """The analysis configurations for the evaluation run.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationRun. + If set, this EvaluationRun will be secured by this key.""" + EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] @@ -4079,6 +4115,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 +4134,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 +4166,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 +4194,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] diff --git a/tests/unit/agentplatform/genai/replays/test_create_evaluation_run.py b/tests/unit/agentplatform/genai/replays/test_create_evaluation_run.py index 738a0b25b2..0b94761fa8 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_evaluation_run.py +++ b/tests/unit/agentplatform/genai/replays/test_create_evaluation_run.py @@ -751,6 +751,31 @@ def test_create_eval_run_with_interactions_data_source(mock_uuid4, client): assert evaluation_run.error is None +_KMS_KEY = ( + "projects/977012026409/locations/us-central1" + "/keyRings/test-kr/cryptoKeys/test-key" +) + + +def test_create_eval_run_with_cmek(client): + """CMEK: encryption_spec is forwarded in the request and returned on the resource.""" + client._api_client._http_options.api_version = "v1beta1" + evaluation_run = client.evals.create_evaluation_run( + name="test_cmek", + display_name="test_cmek", + dataset=types.EvaluationRunDataSource(evaluation_set=EVAL_SET_NAME), + dest=GCS_DEST, + metrics=[GENERAL_QUALITY_METRIC], + encryption_spec=genai_types.EncryptionSpec(kms_key_name=_KMS_KEY), + ) + assert isinstance(evaluation_run, types.EvaluationRun) + assert evaluation_run.display_name == "test_cmek" + assert evaluation_run.state == types.EvaluationRunState.PENDING + assert evaluation_run.encryption_spec is not None + assert evaluation_run.encryption_spec.kms_key_name == _KMS_KEY + assert evaluation_run.error is None + + def test_create_eval_run_with_red_teaming_config(client): """Tests that create_evaluation_run() with red_teaming_config sends analysisConfigs.""" evaluation_run = client.evals.create_evaluation_run( diff --git a/tests/unit/agentplatform/genai/replays/test_create_evaluation_set.py b/tests/unit/agentplatform/genai/replays/test_create_evaluation_set.py index 694c94851f..2a69fd7fec 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_evaluation_set.py +++ b/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",) 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__,