diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index d408ced686..6af15f5fb1 100644 --- a/agentplatform/_genai/evals.py +++ b/agentplatform/_genai/evals.py @@ -187,6 +187,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 @@ -1514,6 +1517,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. @@ -1523,6 +1527,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]] @@ -3527,18 +3532,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. @@ -3547,17 +3554,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, @@ -4164,6 +4175,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. @@ -4173,6 +4185,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]] @@ -5828,6 +5841,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. @@ -5836,6 +5850,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. @@ -5845,6 +5861,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 diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index dbd9c99c90..3080c72882 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -4250,6 +4250,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): @@ -4264,6 +4269,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 @@ -4292,6 +4301,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): @@ -4315,6 +4329,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_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",)