feat: Add extra_context to EvaluationItemRequest

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.")

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]
 

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

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(