diff --git a/agentplatform/_genai/_evals_common.py b/agentplatform/_genai/_evals_common.py index b448395ecb..bd1dac8b81 100644 --- a/agentplatform/_genai/_evals_common.py +++ b/agentplatform/_genai/_evals_common.py @@ -4114,6 +4114,40 @@ def _get_content(row: dict[str, Any], column: str) -> Optional[genai_types.Conte ) +# Columns mapped to EvaluationItemRequest fields; all others go to extra_context. +_EVALUATION_SET_STANDARD_COLUMNS = _evals_constant.COMMON_DATASET_COLUMNS | { + _evals_constant.RESPONSE, + "interaction", + "gemini_agent", +} + + +def _is_missing(value: Any) -> bool: + return value is None or (pd.api.types.is_scalar(value) and pd.isna(value)) + + +def _to_extra_context_content(value: Any) -> Optional[genai_types.Content]: + if isinstance(value, str): + return genai_types.Content(parts=[genai_types.Part(text=value)]) + if isinstance(value, genai_types.Content): + return value + parts = [] + for content in _evals_metric_handlers._value_to_content_list(value): + parts.extend(content.parts or []) + return genai_types.Content(parts=parts) if parts else None + + +def _get_extra_context(row: "pd.Series") -> Optional[dict[str, genai_types.Content]]: + extra_context = {} + for column, value in row.items(): + if column in _EVALUATION_SET_STANDARD_COLUMNS or _is_missing(value): + continue + content = _to_extra_context_content(value) + if content: + extra_context[str(column)] = content + return extra_context or None + + def _create_evaluation_set_from_dataframe( api_client: BaseApiClient, gcs_dest_prefix: str, @@ -4236,6 +4270,7 @@ def _create_evaluation_set_from_dataframe( candidate_responses=( candidate_responses if candidate_responses else None ), + extra_context=_get_extra_context(row), ) ) logger.info("Writing evaluation item requests to GCS.") diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index d59acc7cfa..7706b6b8e7 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -1163,6 +1163,10 @@ class EvaluationItemRequest(_common.BaseModel): default=None, description="""Responses from model under test and other baseline models for comparison.""", ) + extra_context: Optional[dict[str, genai_types.Content]] = Field( + default=None, + description="""Additional named inputs for metrics, such as `golden_schema` alongside an agent trace. Each entry is passed to metrics under its key: LLM-based metrics reference it as `{key}` in the metric prompt template. Standard fields (prompt, response, reference) take precedence over a matching key; `extra_context` takes precedence over `prompt.prompt_template_data.values`.""", + ) class EvaluationItemRequestDict(TypedDict, total=False): @@ -1180,6 +1184,9 @@ class EvaluationItemRequestDict(TypedDict, total=False): candidate_responses: Optional[list[CandidateResponseDict]] """Responses from model under test and other baseline models for comparison.""" + extra_context: Optional[dict[str, genai_types.Content]] + """Additional named inputs for metrics, such as `golden_schema` alongside an agent trace. Each entry is passed to metrics under its key: LLM-based metrics reference it as `{key}` in the metric prompt template. Standard fields (prompt, response, reference) take precedence over a matching key; `extra_context` takes precedence over `prompt.prompt_template_data.values`.""" + EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] 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 395ec37053..0717f17e08 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_evaluation_run.py +++ b/tests/unit/agentplatform/genai/replays/test_create_evaluation_run.py @@ -17,6 +17,7 @@ from tests.unit.agentplatform.genai.replays import pytest_helper from agentplatform import types from agentplatform._genai import _evals_common +from agentplatform._genai import _gcs_utils from google.genai import types as genai_types import pandas as pd import pytest @@ -353,6 +354,65 @@ def test_create_eval_run_with_metric_resource_name( assert evaluation_run.evaluation_config.metrics[0].metric == "my_custom_metric" +# Recorded against staging in ucaip-e2e-test, whose bucket the staging service +# agent can read; staging runs the extra_context backend changes. +@mock.patch.object( + _evals_common, "_local_timestamp", return_value="1/1/2026, 12:00:00 AM" +) +@mock.patch("uuid.uuid4") +def test_create_eval_run_with_extra_context( + mock_uuid4, unused_mock_local_timestamp, client +): + """Tests that custom dataframe columns are sent as extra_context.""" + mock_uuid4.return_value = uuid.UUID("5f0e8a2c-6a1b-4c3d-9e7f-1a2b3c4d5e6f") + client._api_client._http_options.api_version = "v1beta1" + client._api_client._http_options.base_url = ( + "https://us-central1-staging-aiplatform.sandbox.googleapis.com/" + ) + schema_metric = types.EvaluationRunMetric( + metric="schema_check", + metric_config=types.UnifiedMetric( + llm_based_metric_spec=genai_types.LLMBasedMetricSpec( + metric_prompt_template=( + "Schema: {golden_schema}\nResponse: {response}\nScore 1 if the" + " response only uses tables defined in the schema, else 0." + ' Reply with a JSON object with keys "score" and' + ' "explanation".' + ) + ) + ), + ) + with mock.patch.object( + _gcs_utils.GcsUtils, + "upload_json_to_prefix", + autospec=True, + side_effect=_gcs_utils.GcsUtils.upload_json_to_prefix, + ) as mock_upload: + evaluation_run = client.evals.create_evaluation_run( + dataset=types.EvaluationDataset( + eval_dataset_df=pd.DataFrame( + { + "prompt": ["List the ids of all users."], + "response": ["SELECT id FROM users"], + "golden_schema": ["CREATE TABLE users (id INT, name STRING)"], + } + ) + ), + metrics=[schema_metric], + dest="gs://eval-task-e2e-test/sdk_extra_context", + ) + uploaded_request = mock_upload.call_args.kwargs["data"] + assert uploaded_request["extraContext"] == { + "golden_schema": { + "parts": [{"text": "CREATE TABLE users (id INT, name STRING)"}] + } + } + assert isinstance(evaluation_run, types.EvaluationRun) + assert evaluation_run.state == types.EvaluationRunState.PENDING + assert evaluation_run.error is None + assert evaluation_run.evaluation_config.metrics[0].metric == "schema_check" + + # Dataframe tests fail in replay mode because of UUID generation mismatch. # def test_create_eval_run_data_source_evaluation_dataset(client): # """Tests that create_evaluation_run() creates a correctly structured diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py index 28807dfc67..782c0f71d9 100644 --- a/tests/unit/agentplatform/genai/test_evals.py +++ b/tests/unit/agentplatform/genai/test_evals.py @@ -10092,6 +10092,58 @@ def test_convert_request_to_dataset_row_with_agent_data(self): assert result["intermediate_events"] == [] +class TestEvaluationItemRequestExtraContext: + """Unit tests for EvaluationItemRequest.extra_context.""" + + def test_extra_context_serializes_as_camel_case(self): + request = agentplatform_genai_types.EvaluationItemRequest( + prompt=agentplatform_genai_types.EvaluationPrompt(text="test prompt"), + extra_context={ + "golden_schema": genai_types.Content( + parts=[genai_types.Part(text="CREATE TABLE t (id INT)")] + ) + }, + ) + assert request.model_dump(mode="json", by_alias=True, exclude_none=True) == { + "prompt": {"text": "test prompt"}, + "extraContext": { + "golden_schema": {"parts": [{"text": "CREATE TABLE t (id INT)"}]} + }, + } + + def test_extra_context_from_dict(self): + request = agentplatform_genai_types.EvaluationItemRequest.model_validate( + {"extra_context": {"golden_rows": {"parts": [{"text": "[1, 2]"}]}}} + ) + assert request.extra_context == { + "golden_rows": genai_types.Content(parts=[genai_types.Part(text="[1, 2]")]) + } + + @mock.patch.object(_evals_common, "_gcs_utils") + def test_convert_gcs_to_evaluation_item_request_with_extra_context( + self, mock_gcs_utils + ): + """Tests that a camelCase request JSON from GCS keeps extra_context.""" + mock_gcs_utils.GcsUtils.return_value.read_file_contents.return_value = ( + json.dumps( + { + "prompt": {"text": "test prompt"}, + "extraContext": { + "golden_schema": {"parts": [{"text": "CREATE TABLE t"}]} + }, + } + ) + ) + request = _evals_common._convert_gcs_to_evaluation_item_request( + mock.Mock(), "gs://bucket/request.json" + ) + assert request.extra_context == { + "golden_schema": genai_types.Content( + parts=[genai_types.Part(text="CREATE TABLE t")] + ) + } + + class TestCreateEvaluationSetFromDataFrame: """Unit tests for the _create_evaluation_set_from_dataframe function.""" @@ -10100,6 +10152,68 @@ def setup_method(self): self.mock_api_client.project = "test-project" self.mock_api_client.location = "us-central1" + def _uploaded_requests(self, eval_df): + """Returns the request dicts uploaded to GCS for `eval_df`.""" + with ( + mock.patch.object(_evals_common, "evals"), + mock.patch.object(_evals_common, "_gcs_utils") as mock_gcs_utils, + ): + _evals_common._create_evaluation_set_from_dataframe( + api_client=self.mock_api_client, + gcs_dest_prefix="gs://bucket/prefix", + eval_df=eval_df, + ) + upload = mock_gcs_utils.GcsUtils.return_value.upload_json_to_prefix + return [c.kwargs["data"] for c in upload.call_args_list] + + def test_custom_columns_go_to_extra_context(self): + """Tests that only non-standard columns go to extra_context.""" + eval_df = pd.DataFrame( + [ + { + "prompt": "test prompt", + "response": "test response", + "reference": "test reference", + "context": "test context", + "golden_schema": "CREATE TABLE t (id INT)", + "golden_rows": [1, 2], + } + ] + ) + (request,) = self._uploaded_requests(eval_df) + assert request["extraContext"] == { + "golden_schema": {"parts": [{"text": "CREATE TABLE t (id INT)"}]}, + "golden_rows": {"parts": [{"text": "[1, 2]"}]}, + } + + def test_custom_columns_kept_for_agent_data_only_row(self): + eval_df = pd.DataFrame( + [{"agent_data": {"turns": []}, "golden_schema": "CREATE TABLE t"}] + ) + (request,) = self._uploaded_requests(eval_df) + assert "prompt" not in request + assert request["extraContext"] == { + "golden_schema": {"parts": [{"text": "CREATE TABLE t"}]} + } + + def test_missing_custom_values_are_skipped(self): + eval_df = pd.DataFrame( + [ + {"prompt": "p1", "golden_schema": "CREATE TABLE t"}, + {"prompt": "p2", "golden_schema": float("nan")}, + ] + ) + first, second = self._uploaded_requests(eval_df) + assert first["extraContext"] == { + "golden_schema": {"parts": [{"text": "CREATE TABLE t"}]} + } + assert "extraContext" not in second + + def test_no_extra_context_without_custom_columns(self): + eval_df = pd.DataFrame([{"prompt": "test prompt", "response": "test response"}]) + (request,) = self._uploaded_requests(eval_df) + assert "extraContext" not in request + @mock.patch.object(_evals_common, "evals") @mock.patch.object(_evals_common, "_gcs_utils") def test_create_evaluation_set_with_intermediate_events(