From 17bf91ea6c16e8fa158b4c07cddc345f64740358 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Sun, 8 Nov 2026 03:26:46 -0800 Subject: [PATCH] feat: GenAI Client(evals) - Support Harbor datasets, import modes, and waiting in import_evaluation_set PiperOrigin-RevId: 1016689766 --- agentplatform/_genai/evals.py | 408 +++++++++++++++--- agentplatform/_genai/types/__init__.py | 28 ++ agentplatform/_genai/types/common.py | 179 +++++++- .../replays/test_import_evaluation_set.py | 76 ++++ tests/unit/agentplatform/genai/test_evals.py | 334 ++++++++++++++ 5 files changed, 968 insertions(+), 57 deletions(-) diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index ed6ebdaea2..ece68db284 100644 --- a/agentplatform/_genai/evals.py +++ b/agentplatform/_genai/evals.py @@ -30,6 +30,7 @@ from . import _evals_common from . import _evals_utils +from . import _operations_utils from . import _transformers as t from . import types from .types import evals as evals_types @@ -901,6 +902,19 @@ def _GetEvaluationSetParameters_to_vertex( return to_object +def _GetImportEvaluationSetOperationParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + def _ImportEvaluationSetParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -921,6 +935,12 @@ def _ImportEvaluationSetParameters_to_vertex( if getv(from_object, ["cloud_trace_source"]) is not None: setv(to_object, ["cloudTraceSource"], getv(from_object, ["cloud_trace_source"])) + if getv(from_object, ["harbor_dataset"]) is not None: + setv(to_object, ["harborDataset"], getv(from_object, ["harbor_dataset"])) + + if getv(from_object, ["import_mode"]) is not None: + setv(to_object, ["importMode"], getv(from_object, ["import_mode"])) + if getv(from_object, ["config"]) is not None: setv(to_object, ["config"], getv(from_object, ["config"])) @@ -2661,10 +2681,12 @@ def _import_evaluation_set( self, *, evaluation_set: types.EvaluationSetOrDict, - gcs_destination: genai_types.GcsDestinationOrDict, + gcs_destination: Optional[genai_types.GcsDestinationOrDict] = None, gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + harbor_dataset: Optional[types.HarborDatasetOrDict] = None, + import_mode: Optional[types.ImportMode] = None, config: Optional[types.ImportEvaluationSetConfigOrDict] = None, ) -> types.ImportEvaluationSetOperation: """ @@ -2677,6 +2699,8 @@ def _import_evaluation_set( gcs_source=gcs_source, inline_source=inline_source, cloud_trace_source=cloud_trace_source, + harbor_dataset=harbor_dataset, + import_mode=import_mode, config=config, ) @@ -2737,6 +2761,76 @@ def _import_evaluation_set( self._api_client._verify_response(return_value) return return_value + def _get_import_evaluation_set_operation( + self, + *, + operation_name: str, + config: Optional[types.GetImportEvaluationSetOperationConfigOrDict] = None, + ) -> types.ImportEvaluationSetOperation: + parameter_model = types._GetImportEvaluationSetOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetImportEvaluationSetOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ImportEvaluationSetOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + def list_evaluation_experiments( self, *, config: Optional[types.ListEvaluationExperimentsConfigOrDict] = None ) -> types.ListEvaluationExperimentsResponse: @@ -3047,6 +3141,10 @@ def update_evaluation_experiment( self._api_client._verify_response(return_value) return return_value + # Polling for `import_evaluation_set(wait=True)`. + _IMPORT_EVALUATION_SET_POLL_INTERVAL_SECONDS = 5 + _IMPORT_EVALUATION_SET_TIMEOUT_SECONDS = 30 * 60 + def evaluate_instances( self, *, @@ -3982,60 +4080,154 @@ def delete_evaluation_set( name = name.split("/")[-1] self._delete_evaluation_set(name=name, config=config) + @staticmethod + def _validate_import_evaluation_set_sources( + *, + gcs_destination: object, + gcs_source: object, + inline_source: object, + cloud_trace_source: object, + harbor_dataset: object, + ) -> None: + """Checks the source and the destination of an import. + + Shared by the sync and async surfaces. + + Raises: + ValueError: If not exactly one source is set, or if `gcs_destination` + does not match the source. + """ + sources = (gcs_source, inline_source, cloud_trace_source, harbor_dataset) + if sum(source is not None for source in sources) != 1: + raise ValueError( + "Exactly one of gcs_source, inline_source, cloud_trace_source, or" + " harbor_dataset must be provided." + ) + if harbor_dataset is not None and gcs_destination is not None: + raise ValueError("gcs_destination must not be set with harbor_dataset.") + if harbor_dataset is None and gcs_destination is None: + raise ValueError( + "gcs_destination is required unless harbor_dataset is provided." + ) + + @staticmethod + def _check_completed_import_operation( + operation: types.ImportEvaluationSetOperation, + ) -> None: + """Checks the result of a completed import operation. + + Shared by the sync and async surfaces. Logs a warning if some items failed + to import. + + Raises: + RuntimeError: If the operation failed, or completed without an + EvaluationSet. + """ + if operation.error: + raise RuntimeError( + f"Import operation {operation.name} failed: {operation.error}" + ) + response = operation.response + if response is None or not response.evaluation_set: + raise RuntimeError( + "The import operation completed but returned no evaluation set:" + f" {operation!r}" + ) + metadata = operation.metadata or {} + failed_item_count = int(metadata.get("failedItemCount", 0)) + if failed_item_count: + logger.warning( + "Failed to import %d item(s) into %s. Partial failures: %s", + failed_item_count, + response.evaluation_set, + metadata.get("partialFailures"), + ) + def import_evaluation_set( self, *, evaluation_set: types.EvaluationSetOrDict, - gcs_destination: genai_types.GcsDestinationOrDict, + gcs_destination: Optional[genai_types.GcsDestinationOrDict] = None, gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + harbor_dataset: Optional[types.HarborDatasetOrDict] = None, + import_mode: Optional[types.ImportMode] = None, + wait: bool = False, config: Optional[types.ImportEvaluationSetConfigOrDict] = None, ) -> types.ImportEvaluationSetOperation: """Imports data into an EvaluationSet. - This is a long-running operation that imports data from a source (such as - OpenTelemetry traces) into a managed EvaluationSet. + This starts a long-running operation that imports data from a source (such + as OpenTelemetry traces or a Harbor dataset) into a managed EvaluationSet. + Exactly one of ``gcs_source``, ``inline_source``, ``cloud_trace_source``, + or ``harbor_dataset`` must be provided. Args: evaluation_set: The EvaluationSet to create. Used to specify ``display_name`` and ``metadata``. The ``evaluation_items`` field is - ignored and populated by the import process. + ignored and populated by the import process. If ``import_mode`` is + set, ``name`` selects the existing EvaluationSet to import into, in the + form + ``projects/{project}/locations/{location}/evaluationSets/{evaluation_set}``, + and the other fields are ignored. gcs_destination: The Cloud Storage location where the resulting - EvaluationItem payloads will be stored. - gcs_source: The Cloud Storage source of the input data. Exactly one of - ``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be - provided. - inline_source: The inline source for small payloads (< 4MB). Exactly one - of ``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be - provided. + EvaluationItem payloads will be stored. Required for every source + except ``harbor_dataset``, which must not set it. + gcs_source: The Cloud Storage source of the input data. + inline_source: The inline source for small payloads (< 4MB). cloud_trace_source: The Cloud Trace source for loading traces directly - from Cloud Trace. Exactly one of ``gcs_source``, ``inline_source``, or - ``cloud_trace_source`` must be provided. + from Cloud Trace. + harbor_dataset: A Harbor dataset in Cloud Storage. The import creates one + EvaluationItem of type ``HARBOR_TASK`` for each task in the + ``dataset.toml`` file of the dataset. + import_mode: The mode of the import. ``ImportMode.OVERWRITE`` replaces + the items of the EvaluationSet named by ``evaluation_set.name``. It is + supported only for ``harbor_dataset``, into an EvaluationSet that was + imported from a Harbor dataset. + wait: If True, waits until the import finishes and returns the completed + operation. ``operation.response.evaluation_set`` is the name of the + imported EvaluationSet. If False, returns the operation when the import + starts. config: The optional configuration for the import operation. Returns: - The long-running operation for the import. + The long-running operation for the import. If ``wait`` is True, the + operation is done and succeeded. + + Raises: + ValueError: If the sources or ``gcs_destination`` are not valid. + RuntimeError: If ``wait`` is True and the import fails. + TimeoutError: If ``wait`` is True and the import does not finish in 30 + minutes. The error names the operation. """ - if ( - sum( - source is not None - for source in (gcs_source, inline_source, cloud_trace_source) - ) - != 1 - ): - raise ValueError( - "Exactly one of gcs_source, inline_source, or cloud_trace_source" - " must be provided." - ) - return self._import_evaluation_set( + Evals._validate_import_evaluation_set_sources( + gcs_destination=gcs_destination, + gcs_source=gcs_source, + inline_source=inline_source, + cloud_trace_source=cloud_trace_source, + harbor_dataset=harbor_dataset, + ) + operation = self._import_evaluation_set( evaluation_set=evaluation_set, gcs_destination=gcs_destination, gcs_source=gcs_source, inline_source=inline_source, cloud_trace_source=cloud_trace_source, + harbor_dataset=harbor_dataset, + import_mode=import_mode, config=config, ) + if not wait: + return operation + operation = _operations_utils.await_operation( + operation_name=operation.name, + get_operation_fn=self._get_import_evaluation_set_operation, + poll_interval=Evals._IMPORT_EVALUATION_SET_POLL_INTERVAL_SECONDS, + timeout_seconds=Evals._IMPORT_EVALUATION_SET_TIMEOUT_SECONDS, + ) + Evals._check_completed_import_operation(operation) + return operation def generate_conversation_scenarios( self, @@ -5693,10 +5885,12 @@ async def _import_evaluation_set( self, *, evaluation_set: types.EvaluationSetOrDict, - gcs_destination: genai_types.GcsDestinationOrDict, + gcs_destination: Optional[genai_types.GcsDestinationOrDict] = None, gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + harbor_dataset: Optional[types.HarborDatasetOrDict] = None, + import_mode: Optional[types.ImportMode] = None, config: Optional[types.ImportEvaluationSetConfigOrDict] = None, ) -> types.ImportEvaluationSetOperation: """ @@ -5709,6 +5903,8 @@ async def _import_evaluation_set( gcs_source=gcs_source, inline_source=inline_source, cloud_trace_source=cloud_trace_source, + harbor_dataset=harbor_dataset, + import_mode=import_mode, config=config, ) @@ -5771,6 +5967,78 @@ async def _import_evaluation_set( self._api_client._verify_response(return_value) return return_value + async def _get_import_evaluation_set_operation( + self, + *, + operation_name: str, + config: Optional[types.GetImportEvaluationSetOperationConfigOrDict] = None, + ) -> types.ImportEvaluationSetOperation: + parameter_model = types._GetImportEvaluationSetOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetImportEvaluationSetOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ImportEvaluationSetOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + async def list_evaluation_experiments( self, *, config: Optional[types.ListEvaluationExperimentsConfigOrDict] = None ) -> types.ListEvaluationExperimentsResponse: @@ -6633,57 +6901,87 @@ async def import_evaluation_set( self, *, evaluation_set: types.EvaluationSetOrDict, - gcs_destination: genai_types.GcsDestinationOrDict, + gcs_destination: Optional[genai_types.GcsDestinationOrDict] = None, gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + harbor_dataset: Optional[types.HarborDatasetOrDict] = None, + import_mode: Optional[types.ImportMode] = None, + wait: bool = False, config: Optional[types.ImportEvaluationSetConfigOrDict] = None, ) -> types.ImportEvaluationSetOperation: """Imports data into an EvaluationSet. - This is a long-running operation that imports data from a source (such as - OpenTelemetry traces) into a managed EvaluationSet. + This starts a long-running operation that imports data from a source (such + as OpenTelemetry traces or a Harbor dataset) into a managed EvaluationSet. + Exactly one of ``gcs_source``, ``inline_source``, ``cloud_trace_source``, + or ``harbor_dataset`` must be provided. Args: evaluation_set: The EvaluationSet to create. Used to specify ``display_name`` and ``metadata``. The ``evaluation_items`` field is - ignored and populated by the import process. + ignored and populated by the import process. If ``import_mode`` is + set, ``name`` selects the existing EvaluationSet to import into, in the + form + ``projects/{project}/locations/{location}/evaluationSets/{evaluation_set}``, + and the other fields are ignored. gcs_destination: The Cloud Storage location where the resulting - EvaluationItem payloads will be stored. - gcs_source: The Cloud Storage source of the input data. Exactly one of - ``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be - provided. - inline_source: The inline source for small payloads (< 4MB). Exactly one - of ``gcs_source``, ``inline_source``, or ``cloud_trace_source`` must be - provided. + EvaluationItem payloads will be stored. Required for every source + except ``harbor_dataset``, which must not set it. + gcs_source: The Cloud Storage source of the input data. + inline_source: The inline source for small payloads (< 4MB). cloud_trace_source: The Cloud Trace source for loading traces directly - from Cloud Trace. Exactly one of ``gcs_source``, ``inline_source``, or - ``cloud_trace_source`` must be provided. + from Cloud Trace. + harbor_dataset: A Harbor dataset in Cloud Storage. The import creates one + EvaluationItem of type ``HARBOR_TASK`` for each task in the + ``dataset.toml`` file of the dataset. + import_mode: The mode of the import. ``ImportMode.OVERWRITE`` replaces + the items of the EvaluationSet named by ``evaluation_set.name``. It is + supported only for ``harbor_dataset``, into an EvaluationSet that was + imported from a Harbor dataset. + wait: If True, waits until the import finishes and returns the completed + operation. ``operation.response.evaluation_set`` is the name of the + imported EvaluationSet. If False, returns the operation when the import + starts. config: The optional configuration for the import operation. Returns: - The long-running operation for the import. + The long-running operation for the import. If ``wait`` is True, the + operation is done and succeeded. + + Raises: + ValueError: If the sources or ``gcs_destination`` are not valid. + RuntimeError: If ``wait`` is True and the import fails. + TimeoutError: If ``wait`` is True and the import does not finish in 30 + minutes. The error names the operation. """ - if ( - sum( - source is not None - for source in (gcs_source, inline_source, cloud_trace_source) - ) - != 1 - ): - raise ValueError( - "Exactly one of gcs_source, inline_source, or cloud_trace_source" - " must be provided." - ) - result = await self._import_evaluation_set( + Evals._validate_import_evaluation_set_sources( + gcs_destination=gcs_destination, + gcs_source=gcs_source, + inline_source=inline_source, + cloud_trace_source=cloud_trace_source, + harbor_dataset=harbor_dataset, + ) + operation = await self._import_evaluation_set( evaluation_set=evaluation_set, gcs_destination=gcs_destination, gcs_source=gcs_source, inline_source=inline_source, cloud_trace_source=cloud_trace_source, + harbor_dataset=harbor_dataset, + import_mode=import_mode, config=config, ) - return result + if not wait: + return operation + operation = await _operations_utils.await_operation_async( + operation_name=operation.name, + get_operation_fn=self._get_import_evaluation_set_operation, + poll_interval=Evals._IMPORT_EVALUATION_SET_POLL_INTERVAL_SECONDS, + timeout_seconds=Evals._IMPORT_EVALUATION_SET_TIMEOUT_SECONDS, + ) + Evals._check_completed_import_operation(operation) + return operation async def generate_conversation_scenarios( self, diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 7e35512a35..f18535c80f 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -101,6 +101,7 @@ from .common import _GetExampleStoreParameters from .common import _GetExportPublisherModelOperationParameters from .common import _GetGenerateMemoriesOperationParameters +from .common import _GetImportEvaluationSetOperationParameters from .common import _GetImportFilesOperationParameters from .common import _GetMemoryBankOperationParameters from .common import _GetMemoryBankRequestParameters @@ -938,6 +939,9 @@ from .common import GetExportPublisherModelOperationConfig from .common import GetExportPublisherModelOperationConfigDict from .common import GetExportPublisherModelOperationConfigOrDict +from .common import GetImportEvaluationSetOperationConfig +from .common import GetImportEvaluationSetOperationConfigDict +from .common import GetImportEvaluationSetOperationConfigOrDict from .common import GetImportFilesOperationConfig from .common import GetImportFilesOperationConfigDict from .common import GetImportFilesOperationConfigOrDict @@ -1028,6 +1032,12 @@ from .common import GoogleDriveSourceResourceId from .common import GoogleDriveSourceResourceIdDict from .common import GoogleDriveSourceResourceIdOrDict +from .common import HarborDataset +from .common import HarborDatasetDict +from .common import HarborDatasetOrDict +from .common import HarborTask +from .common import HarborTaskDict +from .common import HarborTaskOrDict from .common import IdentityType from .common import Importance from .common import ImportDataFormat @@ -1037,6 +1047,10 @@ from .common import ImportEvaluationSetOperation from .common import ImportEvaluationSetOperationDict from .common import ImportEvaluationSetOperationOrDict +from .common import ImportEvaluationSetResponse +from .common import ImportEvaluationSetResponseDict +from .common import ImportEvaluationSetResponseOrDict +from .common import ImportMode from .common import ImportRagFilesConfig from .common import ImportRagFilesConfigDict from .common import ImportRagFilesConfigOrDict @@ -2438,6 +2452,9 @@ "EvaluationItemResult", "EvaluationItemResultDict", "EvaluationItemResultOrDict", + "HarborTask", + "HarborTaskDict", + "HarborTaskOrDict", "EvaluationItem", "EvaluationItemDict", "EvaluationItemOrDict", @@ -2591,6 +2608,9 @@ "CreateEvaluationSetConfig", "CreateEvaluationSetConfigDict", "CreateEvaluationSetConfigOrDict", + "HarborDataset", + "HarborDatasetDict", + "HarborDatasetOrDict", "EvaluationSet", "EvaluationSetDict", "EvaluationSetOrDict", @@ -2807,9 +2827,15 @@ "ImportEvaluationSetConfig", "ImportEvaluationSetConfigDict", "ImportEvaluationSetConfigOrDict", + "ImportEvaluationSetResponse", + "ImportEvaluationSetResponseDict", + "ImportEvaluationSetResponseOrDict", "ImportEvaluationSetOperation", "ImportEvaluationSetOperationDict", "ImportEvaluationSetOperationOrDict", + "GetImportEvaluationSetOperationConfig", + "GetImportEvaluationSetOperationConfigDict", + "GetImportEvaluationSetOperationConfigOrDict", "ListEvaluationExperimentsConfig", "ListEvaluationExperimentsConfigDict", "ListEvaluationExperimentsConfigOrDict", @@ -4561,6 +4587,7 @@ "SamplingMethod", "EvaluationRunState", "ImportDataFormat", + "ImportMode", "OptimizeTarget", "MemoryMetadataMergeStrategy", "GenerateMemoriesResponseGeneratedMemoryAction", @@ -4609,6 +4636,7 @@ "_GetEvaluationSetParameters", "_GetEvaluationItemParameters", "_ImportEvaluationSetParameters", + "_GetImportEvaluationSetOperationParameters", "_ListEvaluationExperimentsParameters", "_ListEvaluationMetricsParameters", "_ListEvaluationSetsParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 5354c16b5c..4828e6da60 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -704,6 +704,8 @@ class EvaluationItemType(_common.CaseInSensitiveEnum): """The EvaluationItem is a request to evaluate.""" RESULT = "RESULT" """The EvaluationItem is the result of evaluation.""" + HARBOR_TASK = "HARBOR_TASK" + """The EvaluationItem is a Harbor task. Only `import_evaluation_set` creates items of this type, from a Harbor dataset.""" class SamplingMethod(_common.CaseInSensitiveEnum): @@ -745,6 +747,17 @@ class ImportDataFormat(_common.CaseInSensitiveEnum): """JSONL format where each line is a JSON-encoded EvaluationItem.""" +class ImportMode(_common.CaseInSensitiveEnum): + """The mode of an evaluation set import.""" + + IMPORT_MODE_UNSPECIFIED = "IMPORT_MODE_UNSPECIFIED" + """Creates a new EvaluationSet. 'evaluation_set.name' is ignored.""" + OVERWRITE = "OVERWRITE" + """Replaces the items of the EvaluationSet named by + 'evaluation_set.name'. Supported only for 'harbor_dataset', into an + EvaluationSet that was imported from a Harbor dataset.""" + + class OptimizeTarget(_common.CaseInSensitiveEnum): """Specifies the method for calling the optimize_prompt.""" @@ -1252,6 +1265,39 @@ class EvaluationItemResultDict(TypedDict, total=False): EvaluationItemResultOrDict = Union[EvaluationItemResult, EvaluationItemResultDict] +class HarborTask(_common.BaseModel): + """Represents a Harbor task.""" + + dataset: Optional[str] = Field( + default=None, + description="""Required. Harbor dataset name this task belongs to, in `org/name` format without a version, e.g. "my-org/swebench-verified".""", + ) + task: Optional[str] = Field( + default=None, + description="""Required. The Harbor task name, e.g. "django__django-15098".""", + ) + gcs_uri: Optional[str] = Field( + default=None, + description="""The GCS folder containing the Harbor task, e.g. "gs://BUCKET/PATH".""", + ) + + +class HarborTaskDict(TypedDict, total=False): + """Represents a Harbor task.""" + + dataset: Optional[str] + """Required. Harbor dataset name this task belongs to, in `org/name` format without a version, e.g. "my-org/swebench-verified".""" + + task: Optional[str] + """Required. The Harbor task name, e.g. "django__django-15098".""" + + gcs_uri: Optional[str] + """The GCS folder containing the Harbor task, e.g. "gs://BUCKET/PATH".""" + + +HarborTaskOrDict = Union[HarborTask, HarborTaskDict] + + class EvaluationItem(_common.BaseModel): """EvaluationItem is a single evaluation request or result. @@ -1290,6 +1336,10 @@ class EvaluationItem(_common.BaseModel): error: Optional[genai_types.GoogleRpcStatus] = Field( default=None, description="""Error for the evaluation item.""" ) + harbor_task: Optional[HarborTask] = Field( + default=None, + description="""Output only. The Harbor task. Set when `evaluation_item_type` is `HARBOR_TASK`.""", + ) # TODO(b/448806531): Remove all the overridden _from_response methods once the # ticket is resolved and published. @@ -1344,6 +1394,9 @@ class EvaluationItemDict(TypedDict, total=False): error: Optional[genai_types.GoogleRpcStatus] """Error for the evaluation item.""" + harbor_task: Optional[HarborTaskDict] + """Output only. The Harbor task. Set when `evaluation_item_type` is `HARBOR_TASK`.""" + EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] @@ -3958,6 +4011,25 @@ class _CreateEvaluationSetParametersDict(TypedDict, total=False): ] +class HarborDataset(_common.BaseModel): + """Metadata for a Harbor dataset.""" + + gcs_uri: Optional[str] = Field( + default=None, + description="""Optional. The GCS folder containing the Harbor dataset, e.g. "gs://BUCKET/PATH".""", + ) + + +class HarborDatasetDict(TypedDict, total=False): + """Metadata for a Harbor dataset.""" + + gcs_uri: Optional[str] + """Optional. The GCS folder containing the Harbor dataset, e.g. "gs://BUCKET/PATH".""" + + +HarborDatasetOrDict = Union[HarborDataset, HarborDatasetDict] + + class EvaluationSet(_common.BaseModel): """Represents an evaluation set.""" @@ -3985,6 +4057,10 @@ class EvaluationSet(_common.BaseModel): description="""Customer-managed encryption key spec for this EvaluationSet. If set, this EvaluationSet will be secured by this key.""", ) + harbor_dataset: Optional[HarborDataset] = Field( + default=None, + description="""Output only. The Harbor dataset that the items of this evaluation set were imported from.""", + ) class EvaluationSetDict(TypedDict, total=False): @@ -4012,6 +4088,9 @@ class EvaluationSetDict(TypedDict, total=False): """Customer-managed encryption key spec for this EvaluationSet. If set, this EvaluationSet will be secured by this key.""" + harbor_dataset: Optional[HarborDatasetDict] + """Output only. The Harbor dataset that the items of this evaluation set were imported from.""" + EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] @@ -6345,7 +6424,8 @@ class _ImportEvaluationSetParameters(_common.BaseModel): gcs_destination: Optional[genai_types.GcsDestination] = Field( default=None, description="""The Cloud Storage location where the resulting EvaluationItem - payloads will be stored.""", + payloads will be stored. Required for every source except + 'harbor_dataset', which must not set it.""", ) gcs_source: Optional[EvaluationSetGcsSource] = Field( default=None, description="""Google Cloud Storage location.""" @@ -6357,6 +6437,17 @@ class _ImportEvaluationSetParameters(_common.BaseModel): default=None, description="""Source for loading data directly from Cloud Trace.""", ) + harbor_dataset: Optional[HarborDataset] = Field( + default=None, + description="""A Harbor dataset in Cloud Storage. 'gcs_destination' must not be + set with this source.""", + ) + import_mode: Optional[ImportMode] = Field( + default=None, + description="""The mode of the import. If set, 'evaluation_set.name' selects an + existing EvaluationSet to import into, and its other fields are + ignored.""", + ) config: Optional[ImportEvaluationSetConfig] = Field( default=None, description="""""" ) @@ -6372,7 +6463,8 @@ class _ImportEvaluationSetParametersDict(TypedDict, total=False): gcs_destination: Optional[genai_types.GcsDestination] """The Cloud Storage location where the resulting EvaluationItem - payloads will be stored.""" + payloads will be stored. Required for every source except + 'harbor_dataset', which must not set it.""" gcs_source: Optional[EvaluationSetGcsSourceDict] """Google Cloud Storage location.""" @@ -6383,6 +6475,15 @@ class _ImportEvaluationSetParametersDict(TypedDict, total=False): cloud_trace_source: Optional[EvaluationSetCloudTraceSourceDict] """Source for loading data directly from Cloud Trace.""" + harbor_dataset: Optional[HarborDatasetDict] + """A Harbor dataset in Cloud Storage. 'gcs_destination' must not be + set with this source.""" + + import_mode: Optional[ImportMode] + """The mode of the import. If set, 'evaluation_set.name' selects an + existing EvaluationSet to import into, and its other fields are + ignored.""" + config: Optional[ImportEvaluationSetConfigDict] """""" @@ -6392,6 +6493,26 @@ class _ImportEvaluationSetParametersDict(TypedDict, total=False): ] +class ImportEvaluationSetResponse(_common.BaseModel): + """Response of an evaluation set import.""" + + evaluation_set: Optional[str] = Field( + default=None, description="""The resource name of the EvaluationSet.""" + ) + + +class ImportEvaluationSetResponseDict(TypedDict, total=False): + """Response of an evaluation set import.""" + + evaluation_set: Optional[str] + """The resource name of the EvaluationSet.""" + + +ImportEvaluationSetResponseOrDict = Union[ + ImportEvaluationSetResponse, ImportEvaluationSetResponseDict +] + + class ImportEvaluationSetOperation(_common.BaseModel): """Operation for importing an evaluation set.""" @@ -6411,6 +6532,10 @@ class ImportEvaluationSetOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[ImportEvaluationSetResponse] = Field( + default=None, + description="""The result of the import. Set when the operation succeeds.""", + ) class ImportEvaluationSetOperationDict(TypedDict, total=False): @@ -6428,12 +6553,62 @@ class ImportEvaluationSetOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" + response: Optional[ImportEvaluationSetResponseDict] + """The result of the import. Set when the operation succeeds.""" + ImportEvaluationSetOperationOrDict = Union[ ImportEvaluationSetOperation, ImportEvaluationSetOperationDict ] +class GetImportEvaluationSetOperationConfig(_common.BaseModel): + """Config for getting an evaluation set import operation.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetImportEvaluationSetOperationConfigDict(TypedDict, total=False): + """Config for getting an evaluation set import operation.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetImportEvaluationSetOperationConfigOrDict = Union[ + GetImportEvaluationSetOperationConfig, GetImportEvaluationSetOperationConfigDict +] + + +class _GetImportEvaluationSetOperationParameters(_common.BaseModel): + """Parameters for getting an evaluation set import operation.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetImportEvaluationSetOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetImportEvaluationSetOperationParametersDict(TypedDict, total=False): + """Parameters for getting an evaluation set import operation.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetImportEvaluationSetOperationConfigDict] + """Used to override the default configuration.""" + + +_GetImportEvaluationSetOperationParametersOrDict = Union[ + _GetImportEvaluationSetOperationParameters, + _GetImportEvaluationSetOperationParametersDict, +] + + class ListEvaluationExperimentsConfig(_common.BaseModel): """Config for listing evaluation experiments.""" diff --git a/tests/unit/agentplatform/genai/replays/test_import_evaluation_set.py b/tests/unit/agentplatform/genai/replays/test_import_evaluation_set.py index 2f6a1d2919..b69740678f 100644 --- a/tests/unit/agentplatform/genai/replays/test_import_evaluation_set.py +++ b/tests/unit/agentplatform/genai/replays/test_import_evaluation_set.py @@ -21,6 +21,12 @@ _GCS_OUTPUT_PREFIX = "gs://agent-eval-datasets/eval-import-replay/output/" +# The Harbor tests are recorded against autopush, which serves the Harbor import. +_AUTOPUSH_BASE_URL = "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" +_HARBOR_DATASET_GCS_URI = ( + "gs://agent-eval-fishfood/harbor-swebench-verified-mini/import/" + "swebench-verified-mini" +) def test_import_eval_set(client): @@ -40,6 +46,55 @@ def test_import_eval_set(client): assert "/operations/" in operation.name +def test_import_eval_set_from_harbor_dataset(client): + """Tests import_evaluation_set(wait=True) from a Harbor dataset.""" + client._api_client._http_options.base_url = _AUTOPUSH_BASE_URL + + operation = client.evals.import_evaluation_set( + evaluation_set=types.EvaluationSet(display_name="replay-test-harbor-set"), + harbor_dataset=types.HarborDataset(gcs_uri=_HARBOR_DATASET_GCS_URI), + wait=True, + ) + evaluation_set = client.evals.get_evaluation_set( + name=operation.response.evaluation_set + ) + + assert isinstance(operation, types.ImportEvaluationSetOperation) + assert operation.done + assert "/evaluationSets/" in evaluation_set.name + assert evaluation_set.display_name == "replay-test-harbor-set" + assert evaluation_set.harbor_dataset.gcs_uri == _HARBOR_DATASET_GCS_URI + + +def test_import_eval_set_overwrite_harbor_dataset(client): + """Tests import_evaluation_set() with ImportMode.OVERWRITE.""" + client._api_client._http_options.base_url = _AUTOPUSH_BASE_URL + harbor_dataset = types.HarborDataset(gcs_uri=_HARBOR_DATASET_GCS_URI) + operation = client.evals.import_evaluation_set( + evaluation_set=types.EvaluationSet( + display_name="replay-test-harbor-overwrite-set" + ), + harbor_dataset=harbor_dataset, + wait=True, + ) + evaluation_set = client.evals.get_evaluation_set( + name=operation.response.evaluation_set + ) + + overwrite_operation = client.evals.import_evaluation_set( + evaluation_set=types.EvaluationSet(name=evaluation_set.name), + harbor_dataset=harbor_dataset, + import_mode=types.ImportMode.OVERWRITE, + wait=True, + ) + overwritten_set = client.evals.get_evaluation_set( + name=overwrite_operation.response.evaluation_set + ) + + assert overwrite_operation.response.evaluation_set == evaluation_set.name + assert overwritten_set.harbor_dataset.gcs_uri == _HARBOR_DATASET_GCS_URI + + pytest_plugins = ("pytest_asyncio",) @@ -61,6 +116,27 @@ async def test_import_eval_set_async(client): assert "/operations/" in operation.name +@pytest.mark.asyncio +async def test_import_eval_set_from_harbor_dataset_async(client): + """Tests import_evaluation_set(wait=True) on the async client.""" + client._api_client._http_options.base_url = _AUTOPUSH_BASE_URL + + operation = await client.aio.evals.import_evaluation_set( + evaluation_set=types.EvaluationSet(display_name="replay-test-harbor-set"), + harbor_dataset=types.HarborDataset(gcs_uri=_HARBOR_DATASET_GCS_URI), + wait=True, + ) + evaluation_set = await client.aio.evals.get_evaluation_set( + name=operation.response.evaluation_set + ) + + assert isinstance(operation, types.ImportEvaluationSetOperation) + assert operation.done + assert "/evaluationSets/" in evaluation_set.name + assert evaluation_set.display_name == "replay-test-harbor-set" + assert evaluation_set.harbor_dataset.gcs_uri == _HARBOR_DATASET_GCS_URI + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py index fb738cbd51..d49a4dac8e 100644 --- a/tests/unit/agentplatform/genai/test_evals.py +++ b/tests/unit/agentplatform/genai/test_evals.py @@ -64,6 +64,19 @@ _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" +_IMPORT_OPERATION_NAME = "projects/123/locations/us-central1/operations/789" +_IMPORTED_EVALUATION_SET_NAME = "projects/123/locations/us-central1/evaluationSets/456" +_DONE_IMPORT_OPERATION = { + "name": _IMPORT_OPERATION_NAME, + "done": True, + "response": { + "@type": ( + "type.googleapis.com/google.cloud.aiplatform.v1beta1" + ".ImportEvaluationSetResponse" + ), + "evaluationSet": _IMPORTED_EVALUATION_SET_NAME, + }, +} _evals_common = _genai.evals._evals_common @@ -72,6 +85,11 @@ pytestmark = pytest.mark.usefixtures("google_auth_mock") +def _http_responses(*bodies: dict[str, object]) -> list[mock.MagicMock]: + """Returns mock HTTP responses with the given JSON bodies, in order.""" + return [mock.MagicMock(body=json.dumps(body)) for body in bodies] + + class TestDropEmptyColumns: """Unit tests for the _drop_empty_columns function.""" @@ -14612,6 +14630,17 @@ def setup_method(self, method): self.mock_response.body = json.dumps({"name": "operations/123"}) self.mock_api_client.request.return_value = self.mock_response + def _import_from_harbor_dataset(self, evals_module, **kwargs): + return evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="harbor_set" + ), + harbor_dataset=agentplatform_genai_types.HarborDataset( + gcs_uri="gs://bucket/dataset" + ), + **kwargs, + ) + def _import(self, evals_module): return evals_module.import_evaluation_set( evaluation_set=agentplatform_genai_types.EvaluationSet( @@ -14741,3 +14770,308 @@ def test_import_evaluation_set_requires_exactly_one_source(self): gcs_source=gcs_source, cloud_trace_source=cloud_trace_source, ) + + def test_import_evaluation_set_sends_harbor_dataset_without_gcs_destination( + self, + ): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + self._import_from_harbor_dataset(evals_module) + + request_body = self.mock_api_client.request.call_args[0][2] + assert request_body["harborDataset"] == {"gcs_uri": "gs://bucket/dataset"} + assert "gcsDestination" not in request_body + + def test_import_evaluation_set_sends_import_mode_and_target_name(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + name=_IMPORTED_EVALUATION_SET_NAME + ), + harbor_dataset=agentplatform_genai_types.HarborDataset( + gcs_uri="gs://bucket/dataset" + ), + import_mode=agentplatform_genai_types.ImportMode.OVERWRITE, + ) + + request_body = self.mock_api_client.request.call_args[0][2] + assert request_body["importMode"] == "OVERWRITE" + assert request_body["evaluationSet"] == {"name": _IMPORTED_EVALUATION_SET_NAME} + + def test_import_evaluation_set_rejects_gcs_destination_with_harbor_dataset(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + with pytest.raises(ValueError, match="gcs_destination must not be set"): + self._import_from_harbor_dataset( + evals_module, + gcs_destination=genai_types.GcsDestination( + output_uri_prefix="gs://bucket/output/" + ), + ) + + self.mock_api_client.request.assert_not_called() + + def test_import_evaluation_set_requires_gcs_destination_without_harbor_dataset( + self, + ): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + with pytest.raises(ValueError, match="gcs_destination is required"): + evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="imported_set" + ), + cloud_trace_source=agentplatform_genai_types.EvaluationSetCloudTraceSource( + project_id="test-project" + ), + ) + + self.mock_api_client.request.assert_not_called() + + def test_import_evaluation_set_rejects_harbor_dataset_with_another_source(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + with pytest.raises(ValueError, match="Exactly one of"): + self._import_from_harbor_dataset( + evals_module, + cloud_trace_source=agentplatform_genai_types.EvaluationSetCloudTraceSource( + project_id="test-project" + ), + ) + + self.mock_api_client.request.assert_not_called() + + def test_import_evaluation_set_wait_returns_completed_operation(self): + self.mock_api_client.request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + _DONE_IMPORT_OPERATION, + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + operation = self._import_from_harbor_dataset(evals_module, wait=True) + + assert isinstance( + operation, agentplatform_genai_types.ImportEvaluationSetOperation + ) + assert operation.done + assert operation.response.evaluation_set == _IMPORTED_EVALUATION_SET_NAME + + def test_import_evaluation_set_wait_does_not_get_evaluation_set(self): + self.mock_api_client.request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + _DONE_IMPORT_OPERATION, + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + self._import_from_harbor_dataset(evals_module, wait=True) + + requests = [ + call.args[:2] for call in self.mock_api_client.request.call_args_list + ] + assert requests == [ + ("post", "evaluationSets:import"), + ("get", _IMPORT_OPERATION_NAME), + ] + + @mock.patch("time.sleep", return_value=None) + def test_import_evaluation_set_wait_polls_until_operation_is_done(self, mock_sleep): + self.mock_api_client.request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + {"name": _IMPORT_OPERATION_NAME}, + _DONE_IMPORT_OPERATION, + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + operation = self._import_from_harbor_dataset(evals_module, wait=True) + + assert operation.done + mock_sleep.assert_called_once_with(5) + + def test_import_evaluation_set_wait_raises_when_operation_fails(self): + self.mock_api_client.request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + { + "name": _IMPORT_OPERATION_NAME, + "done": True, + "error": {"code": 3, "message": "dataset.toml was not found"}, + }, + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + with pytest.raises( + RuntimeError, match=r"operations/789 failed: .*dataset\.toml was not found" + ): + self._import_from_harbor_dataset(evals_module, wait=True) + + def test_import_evaluation_set_wait_raises_without_evaluation_set(self): + self.mock_api_client.request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + {"name": _IMPORT_OPERATION_NAME, "done": True}, + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + with pytest.raises(RuntimeError, match="returned no evaluation set"): + self._import_from_harbor_dataset(evals_module, wait=True) + + def test_import_evaluation_set_wait_logs_failed_items(self, caplog): + self.mock_api_client.request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + { + "name": _IMPORT_OPERATION_NAME, + "done": True, + "metadata": { + "importedItemCount": "1", + "failedItemCount": "1", + "partialFailures": [ + {"code": 3, "message": "task-2 has no task.toml"} + ], + }, + "response": {"evaluationSet": _IMPORTED_EVALUATION_SET_NAME}, + }, + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + with caplog.at_level("WARNING", logger=evals.logger.name): + self._import_from_harbor_dataset(evals_module, wait=True) + + assert "Failed to import 1 item(s)" in caplog.text + assert "task-2 has no task.toml" in caplog.text + + +class TestAsyncImportEvaluationSet: + + def setup_method(self, method): + self.mock_api_client = mock.MagicMock() + self.mock_api_client.vertexai = True + self.mock_api_client.async_request = mock.AsyncMock() + + @pytest.mark.asyncio + async def test_async_import_evaluation_set_sends_harbor_dataset_and_import_mode( + self, + ): + self.mock_api_client.async_request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME} + ) + async_evals_module = evals.AsyncEvals(api_client_=self.mock_api_client) + + await async_evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + name=_IMPORTED_EVALUATION_SET_NAME + ), + harbor_dataset=agentplatform_genai_types.HarborDataset( + gcs_uri="gs://bucket/dataset" + ), + import_mode=agentplatform_genai_types.ImportMode.OVERWRITE, + ) + + request_body = self.mock_api_client.async_request.call_args[0][2] + assert request_body["harborDataset"] == {"gcs_uri": "gs://bucket/dataset"} + assert request_body["importMode"] == "OVERWRITE" + + @pytest.mark.asyncio + async def test_async_import_evaluation_set_wait_returns_completed_operation( + self, + ): + self.mock_api_client.async_request.side_effect = _http_responses( + {"name": _IMPORT_OPERATION_NAME}, + _DONE_IMPORT_OPERATION, + ) + async_evals_module = evals.AsyncEvals(api_client_=self.mock_api_client) + + operation = await async_evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="harbor_set" + ), + harbor_dataset=agentplatform_genai_types.HarborDataset( + gcs_uri="gs://bucket/dataset" + ), + wait=True, + ) + + assert operation.done + assert operation.response.evaluation_set == _IMPORTED_EVALUATION_SET_NAME + assert self.mock_api_client.async_request.call_count == 2 + + @pytest.mark.asyncio + async def test_async_import_evaluation_set_rejects_gcs_destination_with_harbor_dataset( + self, + ): + async_evals_module = evals.AsyncEvals(api_client_=self.mock_api_client) + + with pytest.raises(ValueError, match="gcs_destination must not be set"): + await async_evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="harbor_set" + ), + gcs_destination=genai_types.GcsDestination( + output_uri_prefix="gs://bucket/output/" + ), + harbor_dataset=agentplatform_genai_types.HarborDataset( + gcs_uri="gs://bucket/dataset" + ), + ) + + self.mock_api_client.async_request.assert_not_called() + + +class TestGetEvaluationSet: + + def setup_method(self, method): + self.mock_api_client = mock.MagicMock() + self.mock_api_client.vertexai = True + self.mock_response = mock.MagicMock() + self.mock_api_client.request.return_value = self.mock_response + + def test_get_evaluation_set_parses_harbor_dataset(self): + self.mock_response.body = json.dumps( + { + "name": "projects/123/locations/us-central1/evaluationSets/456", + "harborDataset": {"gcsUri": "gs://bucket/dataset"}, + } + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + evaluation_set = evals_module.get_evaluation_set( + name="projects/123/locations/us-central1/evaluationSets/456" + ) + + assert evaluation_set.harbor_dataset == ( + agentplatform_genai_types.HarborDataset(gcs_uri="gs://bucket/dataset") + ) + + +class TestGetEvaluationItem: + + def setup_method(self, method): + self.mock_api_client = mock.MagicMock() + self.mock_api_client.vertexai = True + self.mock_response = mock.MagicMock() + self.mock_api_client.request.return_value = self.mock_response + + def test_get_evaluation_item_parses_harbor_task(self): + self.mock_response.body = json.dumps( + { + "name": "projects/123/locations/us-central1/evaluationItems/789", + "evaluationItemType": "HARBOR_TASK", + "harborTask": { + "dataset": "my-org/my-dataset", + "task": "task-1", + "gcsUri": "gs://bucket/dataset/task-1", + }, + } + ) + evals_module = evals.Evals(api_client_=self.mock_api_client) + + evaluation_item = evals_module.get_evaluation_item( + name="projects/123/locations/us-central1/evaluationItems/789" + ) + + assert evaluation_item.evaluation_item_type == ( + agentplatform_genai_types.EvaluationItemType.HARBOR_TASK + ) + assert evaluation_item.harbor_task == agentplatform_genai_types.HarborTask( + dataset="my-org/my-dataset", + task="task-1", + gcs_uri="gs://bucket/dataset/task-1", + )