From 87358fd53d5e9e16261ce0820a8e3ca41a524481 Mon Sep 17 00:00:00 2001 From: Jason Dai Date: Fri, 7 Aug 2026 18:54:31 -0700 Subject: [PATCH] feat: GenAI Client(evals) - add import_evaluation_set PiperOrigin-RevId: 961225038 --- agentplatform/_genai/evals.py | 299 ++++++++++++++++++ agentplatform/_genai/types/__init__.py | 40 +++ agentplatform/_genai/types/common.py | 239 ++++++++++++++ .../replays/test_import_evaluation_set.py | 68 ++++ tests/unit/agentplatform/genai/test_evals.py | 140 ++++++++ 5 files changed, 786 insertions(+) create mode 100644 tests/unit/agentplatform/genai/replays/test_import_evaluation_set.py diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index 01917a6a55..cc77b0ea6e 100644 --- a/agentplatform/_genai/evals.py +++ b/agentplatform/_genai/evals.py @@ -880,6 +880,32 @@ def _GetEvaluationSetParameters_to_vertex( return to_object +def _ImportEvaluationSetParameters_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, ["evaluation_set"]) is not None: + setv(to_object, ["evaluationSet"], getv(from_object, ["evaluation_set"])) + + if getv(from_object, ["gcs_destination"]) is not None: + setv(to_object, ["gcsDestination"], getv(from_object, ["gcs_destination"])) + + if getv(from_object, ["gcs_source"]) is not None: + setv(to_object, ["gcsSource"], getv(from_object, ["gcs_source"])) + + if getv(from_object, ["inline_source"]) is not None: + setv(to_object, ["inlineSource"], getv(from_object, ["inline_source"])) + + if getv(from_object, ["cloud_trace_source"]) is not None: + setv(to_object, ["cloudTraceSource"], getv(from_object, ["cloud_trace_source"])) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + def _ListEvaluationExperimentsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -2594,6 +2620,86 @@ def _get_evaluation_item( self._api_client._verify_response(return_value) return return_value + def _import_evaluation_set( + self, + *, + evaluation_set: types.EvaluationSetOrDict, + gcs_destination: genai_types.GcsDestinationOrDict, + gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, + inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, + cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + config: Optional[types.ImportEvaluationSetConfigOrDict] = None, + ) -> types.ImportEvaluationSetOperation: + """ + Imports data into an EvaluationSet. + """ + + parameter_model = types._ImportEvaluationSetParameters( + evaluation_set=evaluation_set, + gcs_destination=gcs_destination, + gcs_source=gcs_source, + inline_source=inline_source, + cloud_trace_source=cloud_trace_source, + 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 = _ImportEvaluationSetParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "evaluationSets:import".format_map(request_url_dict) + else: + path = "evaluationSets:import" + + 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("post", 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: @@ -3836,6 +3942,61 @@ def delete_evaluation_set( name = name.split("/")[-1] self._delete_evaluation_set(name=name, config=config) + def import_evaluation_set( + self, + *, + evaluation_set: types.EvaluationSetOrDict, + gcs_destination: genai_types.GcsDestinationOrDict, + gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, + inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, + cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + 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. + + 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. + 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. + 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. + config: The optional configuration for the import operation. + + Returns: + The long-running operation for the import. + """ + 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( + evaluation_set=evaluation_set, + gcs_destination=gcs_destination, + gcs_source=gcs_source, + inline_source=inline_source, + cloud_trace_source=cloud_trace_source, + config=config, + ) + def generate_conversation_scenarios( self, *, @@ -5486,6 +5647,88 @@ async def _get_evaluation_item( self._api_client._verify_response(return_value) return return_value + async def _import_evaluation_set( + self, + *, + evaluation_set: types.EvaluationSetOrDict, + gcs_destination: genai_types.GcsDestinationOrDict, + gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, + inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, + cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + config: Optional[types.ImportEvaluationSetConfigOrDict] = None, + ) -> types.ImportEvaluationSetOperation: + """ + Imports data into an EvaluationSet. + """ + + parameter_model = types._ImportEvaluationSetParameters( + evaluation_set=evaluation_set, + gcs_destination=gcs_destination, + gcs_source=gcs_source, + inline_source=inline_source, + cloud_trace_source=cloud_trace_source, + 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 = _ImportEvaluationSetParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "evaluationSets:import".format_map(request_url_dict) + else: + path = "evaluationSets:import" + + 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( + "post", 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: @@ -6344,6 +6587,62 @@ async def delete_evaluation_set( name = name.split("/")[-1] await self._delete_evaluation_set(name=name, config=config) + async def import_evaluation_set( + self, + *, + evaluation_set: types.EvaluationSetOrDict, + gcs_destination: genai_types.GcsDestinationOrDict, + gcs_source: Optional[types.EvaluationSetGcsSourceOrDict] = None, + inline_source: Optional[types.EvaluationSetInlineSourceOrDict] = None, + cloud_trace_source: Optional[types.EvaluationSetCloudTraceSourceOrDict] = None, + 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. + + 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. + 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. + 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. + config: The optional configuration for the import operation. + + Returns: + The long-running operation for the import. + """ + 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( + evaluation_set=evaluation_set, + gcs_destination=gcs_destination, + gcs_source=gcs_source, + inline_source=inline_source, + cloud_trace_source=cloud_trace_source, + config=config, + ) + return result + async def generate_conversation_scenarios( self, *, diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 45b472f3c5..38dd405f50 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -122,6 +122,7 @@ from .common import _GetSkillOperationParameters from .common import _GetSkillRequestParameters from .common import _GetSkillRevisionRequestParameters +from .common import _ImportEvaluationSetParameters from .common import _ImportRagFilesRequestParameters from .common import _IngestEventsRequestParameters from .common import _ListAgentEngineMemoryRequestParameters @@ -657,7 +658,16 @@ from .common import EvaluationRunResultsOrDict from .common import EvaluationRunState from .common import EvaluationSet +from .common import EvaluationSetCloudTraceSource +from .common import EvaluationSetCloudTraceSourceDict +from .common import EvaluationSetCloudTraceSourceOrDict from .common import EvaluationSetDict +from .common import EvaluationSetGcsSource +from .common import EvaluationSetGcsSourceDict +from .common import EvaluationSetGcsSourceOrDict +from .common import EvaluationSetInlineSource +from .common import EvaluationSetInlineSourceDict +from .common import EvaluationSetInlineSourceOrDict from .common import EvaluationSetOrDict from .common import Event from .common import EventActions @@ -930,6 +940,13 @@ from .common import GoogleDriveSourceResourceIdOrDict from .common import IdentityType from .common import Importance +from .common import ImportDataFormat +from .common import ImportEvaluationSetConfig +from .common import ImportEvaluationSetConfigDict +from .common import ImportEvaluationSetConfigOrDict +from .common import ImportEvaluationSetOperation +from .common import ImportEvaluationSetOperationDict +from .common import ImportEvaluationSetOperationOrDict from .common import ImportRagFilesConfig from .common import ImportRagFilesConfigDict from .common import ImportRagFilesConfigOrDict @@ -945,6 +962,9 @@ from .common import ImportRagFilesResponse from .common import ImportRagFilesResponseDict from .common import ImportRagFilesResponseOrDict +from .common import ImportSchemaConfig +from .common import ImportSchemaConfigDict +from .common import ImportSchemaConfigOrDict from .common import InferenceEventLoggingConfig from .common import InferenceEventLoggingConfigDict from .common import InferenceEventLoggingConfigOrDict @@ -2599,6 +2619,24 @@ "GetEvaluationItemConfig", "GetEvaluationItemConfigDict", "GetEvaluationItemConfigOrDict", + "ImportSchemaConfig", + "ImportSchemaConfigDict", + "ImportSchemaConfigOrDict", + "EvaluationSetGcsSource", + "EvaluationSetGcsSourceDict", + "EvaluationSetGcsSourceOrDict", + "EvaluationSetInlineSource", + "EvaluationSetInlineSourceDict", + "EvaluationSetInlineSourceOrDict", + "EvaluationSetCloudTraceSource", + "EvaluationSetCloudTraceSourceDict", + "EvaluationSetCloudTraceSourceOrDict", + "ImportEvaluationSetConfig", + "ImportEvaluationSetConfigDict", + "ImportEvaluationSetConfigOrDict", + "ImportEvaluationSetOperation", + "ImportEvaluationSetOperationDict", + "ImportEvaluationSetOperationOrDict", "ListEvaluationExperimentsConfig", "ListEvaluationExperimentsConfigDict", "ListEvaluationExperimentsConfigOrDict", @@ -4095,6 +4133,7 @@ "EvaluationItemType", "SamplingMethod", "EvaluationRunState", + "ImportDataFormat", "OptimizeTarget", "MemoryMetadataMergeStrategy", "GenerateMemoriesResponseGeneratedMemoryAction", @@ -4147,6 +4186,7 @@ "_GetEvaluationRunParameters", "_GetEvaluationSetParameters", "_GetEvaluationItemParameters", + "_ImportEvaluationSetParameters", "_ListEvaluationExperimentsParameters", "_ListEvaluationMetricsParameters", "_ListEvaluationSetsParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index f2318e1fe6..30c8a7e07d 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -682,6 +682,15 @@ class EvaluationRunState(_common.CaseInSensitiveEnum): """Evaluation run is performing rubric generation.""" +class ImportDataFormat(_common.CaseInSensitiveEnum): + """The format of the input data for an evaluation set import.""" + + DATA_FORMAT_UNSPECIFIED = "DATA_FORMAT_UNSPECIFIED" + """Unspecified data format.""" + JSONL = "JSONL" + """JSONL format where each line is a JSON-encoded EvaluationItem.""" + + class OptimizeTarget(_common.CaseInSensitiveEnum): """Specifies the method for calling the optimize_prompt.""" @@ -6785,6 +6794,236 @@ class _GetEvaluationItemParametersDict(TypedDict, total=False): ] +class ImportSchemaConfig(_common.BaseModel): + """Configuration for the input data format.""" + + data_format: Optional[ImportDataFormat] = Field( + default=None, description="""The format of the input data.""" + ) + data_format_version: Optional[str] = Field( + default=None, description="""Version of the data format.""" + ) + + +class ImportSchemaConfigDict(TypedDict, total=False): + """Configuration for the input data format.""" + + data_format: Optional[ImportDataFormat] + """The format of the input data.""" + + data_format_version: Optional[str] + """Version of the data format.""" + + +ImportSchemaConfigOrDict = Union[ImportSchemaConfig, ImportSchemaConfigDict] + + +class EvaluationSetGcsSource(_common.BaseModel): + """Source for loading data from Cloud Storage.""" + + gcs_uri: Optional[str] = Field( + default=None, description="""The Cloud Storage location of the input data.""" + ) + import_schema_config: Optional[ImportSchemaConfig] = Field( + default=None, description="""Schema configuration for the input data.""" + ) + + +class EvaluationSetGcsSourceDict(TypedDict, total=False): + """Source for loading data from Cloud Storage.""" + + gcs_uri: Optional[str] + """The Cloud Storage location of the input data.""" + + import_schema_config: Optional[ImportSchemaConfigDict] + """Schema configuration for the input data.""" + + +EvaluationSetGcsSourceOrDict = Union[EvaluationSetGcsSource, EvaluationSetGcsSourceDict] + + +class EvaluationSetInlineSource(_common.BaseModel): + """Wrapper for inline data.""" + + content: Optional[bytes] = Field( + default=None, description="""The content of the inline data.""" + ) + import_schema_config: Optional[ImportSchemaConfig] = Field( + default=None, description="""Schema configuration for the inline data.""" + ) + + +class EvaluationSetInlineSourceDict(TypedDict, total=False): + """Wrapper for inline data.""" + + content: Optional[bytes] + """The content of the inline data.""" + + import_schema_config: Optional[ImportSchemaConfigDict] + """Schema configuration for the inline data.""" + + +EvaluationSetInlineSourceOrDict = Union[ + EvaluationSetInlineSource, EvaluationSetInlineSourceDict +] + + +class EvaluationSetCloudTraceSource(_common.BaseModel): + """Source for loading traces directly from Cloud Trace.""" + + project_id: Optional[str] = Field( + default=None, description="""Project ID for the Cloud Trace.""" + ) + trace_ids: Optional[list[str]] = Field( + default=None, description="""Trace IDs to import.""" + ) + session_ids: Optional[list[str]] = Field( + default=None, + description="""Session IDs to import traces for. If both trace_ids and + session_ids are specified, the union of the two will be imported.""", + ) + + +class EvaluationSetCloudTraceSourceDict(TypedDict, total=False): + """Source for loading traces directly from Cloud Trace.""" + + project_id: Optional[str] + """Project ID for the Cloud Trace.""" + + trace_ids: Optional[list[str]] + """Trace IDs to import.""" + + session_ids: Optional[list[str]] + """Session IDs to import traces for. If both trace_ids and + session_ids are specified, the union of the two will be imported.""" + + +EvaluationSetCloudTraceSourceOrDict = Union[ + EvaluationSetCloudTraceSource, EvaluationSetCloudTraceSourceDict +] + + +class ImportEvaluationSetConfig(_common.BaseModel): + """Config for importing an evaluation set.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class ImportEvaluationSetConfigDict(TypedDict, total=False): + """Config for importing an evaluation set.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +ImportEvaluationSetConfigOrDict = Union[ + ImportEvaluationSetConfig, ImportEvaluationSetConfigDict +] + + +class _ImportEvaluationSetParameters(_common.BaseModel): + """Parameters for importing an evaluation set.""" + + evaluation_set: Optional[EvaluationSet] = Field( + default=None, + description="""The EvaluationSet to create. Used to specify 'display_name' and + 'metadata'. The 'evaluation_items' field is ignored and populated by the + import process.""", + ) + gcs_destination: Optional[genai_types.GcsDestination] = Field( + default=None, + description="""The Cloud Storage location where the resulting EvaluationItem + payloads will be stored.""", + ) + gcs_source: Optional[EvaluationSetGcsSource] = Field( + default=None, description="""Google Cloud Storage location.""" + ) + inline_source: Optional[EvaluationSetInlineSource] = Field( + default=None, description="""Inline source for small payloads (< 4MB).""" + ) + cloud_trace_source: Optional[EvaluationSetCloudTraceSource] = Field( + default=None, + description="""Source for loading data directly from Cloud Trace.""", + ) + config: Optional[ImportEvaluationSetConfig] = Field( + default=None, description="""""" + ) + + +class _ImportEvaluationSetParametersDict(TypedDict, total=False): + """Parameters for importing an evaluation set.""" + + evaluation_set: Optional[EvaluationSetDict] + """The EvaluationSet to create. Used to specify 'display_name' and + 'metadata'. The 'evaluation_items' field is ignored and populated by the + import process.""" + + gcs_destination: Optional[genai_types.GcsDestination] + """The Cloud Storage location where the resulting EvaluationItem + payloads will be stored.""" + + gcs_source: Optional[EvaluationSetGcsSourceDict] + """Google Cloud Storage location.""" + + inline_source: Optional[EvaluationSetInlineSourceDict] + """Inline source for small payloads (< 4MB).""" + + cloud_trace_source: Optional[EvaluationSetCloudTraceSourceDict] + """Source for loading data directly from Cloud Trace.""" + + config: Optional[ImportEvaluationSetConfigDict] + """""" + + +_ImportEvaluationSetParametersOrDict = Union[ + _ImportEvaluationSetParameters, _ImportEvaluationSetParametersDict +] + + +class ImportEvaluationSetOperation(_common.BaseModel): + """Operation for importing an evaluation set.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class ImportEvaluationSetOperationDict(TypedDict, total=False): + """Operation for importing an evaluation set.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +ImportEvaluationSetOperationOrDict = Union[ + ImportEvaluationSetOperation, ImportEvaluationSetOperationDict +] + + 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 new file mode 100644 index 0000000000..2f6a1d2919 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_import_evaluation_set.py @@ -0,0 +1,68 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.agentplatform.genai.replays import pytest_helper +from agentplatform import types +from google.genai import types as genai_types +import pytest + + +_GCS_OUTPUT_PREFIX = "gs://agent-eval-datasets/eval-import-replay/output/" + + +def test_import_eval_set(client): + """Tests import_evaluation_set() from a Cloud Trace source.""" + operation = client.evals.import_evaluation_set( + evaluation_set=types.EvaluationSet(display_name="replay-test-import-set"), + gcs_destination=genai_types.GcsDestination( + output_uri_prefix=_GCS_OUTPUT_PREFIX + ), + cloud_trace_source=types.EvaluationSetCloudTraceSource( + project_id="vertex-sdk-dev", + session_ids=["replay-session-1"], + ), + ) + assert isinstance(operation, types.ImportEvaluationSetOperation) + assert operation.name is not None + assert "/operations/" in operation.name + + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_import_eval_set_async(client): + """Tests import_evaluation_set() on the async client.""" + operation = await client.aio.evals.import_evaluation_set( + evaluation_set=types.EvaluationSet(display_name="replay-test-import-set"), + gcs_destination=genai_types.GcsDestination( + output_uri_prefix=_GCS_OUTPUT_PREFIX + ), + cloud_trace_source=types.EvaluationSetCloudTraceSource( + project_id="vertex-sdk-dev", + session_ids=["replay-session-1"], + ), + ) + assert isinstance(operation, types.ImportEvaluationSetOperation) + assert operation.name is not None + assert "/operations/" in operation.name + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), + test_method="evals.import_evaluation_set", +) diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py index 83272b64fb..1e7670bd2f 100644 --- a/tests/unit/agentplatform/genai/test_evals.py +++ b/tests/unit/agentplatform/genai/test_evals.py @@ -12700,3 +12700,143 @@ async def test_async_uses_provided_experiment_without_creating(self): request_body.get("evaluationExperiment") == "projects/123/locations/us-central1/evaluationExperiments/existing" ) + + +class TestImportEvaluationSet: + + def setup_method(self, method): + self.mock_api_client = mock.MagicMock() + self.mock_api_client.vertexai = True + self.mock_response = mock.MagicMock() + self.mock_response.body = json.dumps({"name": "operations/123"}) + self.mock_api_client.request.return_value = self.mock_response + + def _import(self, evals_module): + return evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="imported_set" + ), + gcs_destination=genai_types.GcsDestination( + output_uri_prefix="gs://bucket/output/" + ), + gcs_source=agentplatform_genai_types.EvaluationSetGcsSource( + gcs_uri="gs://bucket/items.jsonl", + import_schema_config=agentplatform_genai_types.ImportSchemaConfig( + data_format=agentplatform_genai_types.ImportDataFormat.JSONL + ), + ), + ) + + def test_import_evaluation_set_posts_to_import_path(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + self._import(evals_module) + + self.mock_api_client.request.assert_called_once() + call_args = self.mock_api_client.request.call_args + assert call_args[0][0] == "post" + assert call_args[0][1] == "evaluationSets:import" + + def test_import_evaluation_set_sends_gcs_source(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + self._import(evals_module) + + request_body = self.mock_api_client.request.call_args[0][2] + gcs_source = request_body["gcsSource"] + assert gcs_source["gcs_uri"] == "gs://bucket/items.jsonl" + assert gcs_source["import_schema_config"]["data_format"] == "JSONL" + + def test_import_evaluation_set_sends_gcs_destination(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + self._import(evals_module) + + request_body = self.mock_api_client.request.call_args[0][2] + assert ( + request_body["gcsDestination"]["output_uri_prefix"] == "gs://bucket/output/" + ) + + def test_import_evaluation_set_returns_operation(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + operation = self._import(evals_module) + + assert operation.name == "operations/123" + + def test_import_evaluation_set_from_inline_source(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="imported_set" + ), + gcs_destination=genai_types.GcsDestination( + output_uri_prefix="gs://bucket/output/" + ), + inline_source=agentplatform_genai_types.EvaluationSetInlineSource( + content=b'{"evaluationRequest": {}}\n', + import_schema_config=agentplatform_genai_types.ImportSchemaConfig( + data_format=agentplatform_genai_types.ImportDataFormat.JSONL + ), + ), + ) + + request_body = self.mock_api_client.request.call_args[0][2] + assert "inlineSource" in request_body + assert ( + request_body["inlineSource"]["import_schema_config"]["data_format"] + == "JSONL" + ) + + def test_import_evaluation_set_from_cloud_trace_source(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + + evals_module.import_evaluation_set( + evaluation_set=agentplatform_genai_types.EvaluationSet( + display_name="imported_set" + ), + gcs_destination=genai_types.GcsDestination( + output_uri_prefix="gs://bucket/output/" + ), + cloud_trace_source=agentplatform_genai_types.EvaluationSetCloudTraceSource( + project_id="test-project", + session_ids=["session-1"], + ), + ) + + request_body = self.mock_api_client.request.call_args[0][2] + cloud_trace_source = request_body["cloudTraceSource"] + assert cloud_trace_source["project_id"] == "test-project" + assert cloud_trace_source["session_ids"] == ["session-1"] + + def test_import_evaluation_set_requires_exactly_one_source(self): + evals_module = evals.Evals(api_client_=self.mock_api_client) + evaluation_set = agentplatform_genai_types.EvaluationSet( + display_name="imported_set" + ) + gcs_destination = genai_types.GcsDestination( + output_uri_prefix="gs://bucket/output/" + ) + gcs_source = agentplatform_genai_types.EvaluationSetGcsSource( + gcs_uri="gs://bucket/items.jsonl", + import_schema_config=agentplatform_genai_types.ImportSchemaConfig( + data_format=agentplatform_genai_types.ImportDataFormat.JSONL + ), + ) + cloud_trace_source = agentplatform_genai_types.EvaluationSetCloudTraceSource( + project_id="test-project" + ) + + with pytest.raises(ValueError): + evals_module.import_evaluation_set( + evaluation_set=evaluation_set, gcs_destination=gcs_destination + ) + + with pytest.raises(ValueError): + evals_module.import_evaluation_set( + evaluation_set=evaluation_set, + gcs_destination=gcs_destination, + gcs_source=gcs_source, + cloud_trace_source=cloud_trace_source, + )