chore: add evals data converter to vertex_genai

PiperOrigin-RevId: 768538090

Author: jsondai

tests/unit/vertexai/genai/test_evals.py

@@ -1208,3 +1208,147 @@ def test_convert_multiple_items(self):
         assert len(result_dataset.eval_cases) == 2
         assert result_dataset.eval_cases[0].prompt.parts[0].text == "Item 1"
         assert result_dataset.eval_cases[1].prompt.parts[0].text == "Item 2"
+
+
+class TestFlattenEvalDataConverter:
+    """Unit tests for the _FlattenEvalDataConverter class."""
+
+    def setup_method(self):
+        self.converter = _evals_data_converters._FlattenEvalDataConverter()
+
+    def test_convert_simple_prompt_response(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": ["Hello"],
+                "response": ["Hi"],
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        result_dataset = self.converter.convert(raw_data)
+        assert isinstance(result_dataset, vertexai_genai_types.EvaluationDataset)
+        assert len(result_dataset.eval_cases) == 1
+        eval_case = result_dataset.eval_cases[0]
+
+        assert eval_case.prompt == genai_types.Content(
+            parts=[genai_types.Part(text="Hello")]
+        )
+        assert len(eval_case.responses) == 1
+        assert eval_case.responses[0].response == genai_types.Content(
+            parts=[genai_types.Part(text="Hi")]
+        )
+        assert eval_case.reference is None
+        assert eval_case.system_instruction is None
+        assert eval_case.conversation_history is None
+
+    def test_convert_with_system_instruction_and_reference(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": ["Hello"],
+                "response": ["Hi there!"],
+                "instruction": ["Be nice."],
+                "reference": ["Hey"],
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        result_dataset = self.converter.convert(raw_data)
+        eval_case = result_dataset.eval_cases[0]
+        assert eval_case.system_instruction == genai_types.Content(
+            parts=[genai_types.Part(text="Be nice.")]
+        )
+        assert eval_case.prompt == genai_types.Content(
+            parts=[genai_types.Part(text="Hello")]
+        )
+        assert eval_case.reference.response == genai_types.Content(
+            parts=[genai_types.Part(text="Hey")]
+        )
+
+    def test_convert_with_conversation_history(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": ["Current prompt"],
+                "response": ["A response"],
+                "history": [
+                    [
+                        {"role": "user", "parts": [{"text": "Old user msg"}]},
+                        {"role": "model", "parts": [{"text": "Old model msg"}]},
+                    ]
+                ],
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        result_dataset = self.converter.convert(raw_data)
+        eval_case = result_dataset.eval_cases[0]
+
+        assert eval_case.prompt == genai_types.Content(
+            parts=[genai_types.Part(text="Current prompt")]
+        )
+        assert eval_case.reference is None
+        assert len(eval_case.conversation_history) == 2
+        assert eval_case.conversation_history[0].content.parts[0].text == "Old user msg"
+        assert (
+            eval_case.conversation_history[1].content.parts[0].text == "Old model msg"
+        )
+
+    def test_convert_missing_response_raises_value_error(self):
+        raw_data_df = pd.DataFrame({"prompt": ["Hello"]})  # Missing response
+        raw_data = raw_data_df.to_dict(orient="records")
+        with pytest.raises(
+            ValueError, match="Response is required but missing for eval_case_0"
+        ):
+            self.converter.convert(raw_data)
+
+    def test_convert_missing_prompt_raises_value_error(self):
+        raw_data_df = pd.DataFrame({"response": ["Hi"]})  # Missing prompt
+        raw_data = raw_data_df.to_dict(orient="records")
+        with pytest.raises(
+            ValueError, match="Prompt is required but missing for eval_case_0"
+        ):
+            self.converter.convert(raw_data)
+
+    def test_convert_invalid_prompt_type_raises_value_error(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": [123],  # Invalid prompt type
+                "response": ["Hi"],
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        with pytest.raises(ValueError, match="Invalid prompt type for case 0"):
+            self.converter.convert(raw_data)
+
+    def test_convert_invalid_response_type_raises_value_error(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": ["Hello"],
+                "response": [123],  # Invalid response type
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        with pytest.raises(ValueError, match="Invalid response type for case 0"):
+            self.converter.convert(raw_data)
+
+    def test_convert_multiple_items(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": ["Item 1", "Item 2"],
+                "response": ["Resp 1", "Resp 2"],
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        result_dataset = self.converter.convert(raw_data)
+        assert len(result_dataset.eval_cases) == 2
+        assert result_dataset.eval_cases[0].prompt.parts[0].text == "Item 1"
+        assert result_dataset.eval_cases[1].prompt.parts[0].text == "Item 2"
+
+    def test_convert_with_additional_columns(self):
+        raw_data_df = pd.DataFrame(
+            {
+                "prompt": ["Hello"],
+                "response": ["Hi"],
+                "custom_column": ["custom_value"],
+            }
+        )
+        raw_data = raw_data_df.to_dict(orient="records")
+        result_dataset = self.converter.convert(raw_data)
+        eval_case = result_dataset.eval_cases[0]
+        assert eval_case.custom_column == "custom_value"

vertexai/_genai/_evals_data_converters.py

@@ -16,7 +16,7 @@
 import abc
 import enum
 import logging
-from typing import Any
+from typing import Any, Optional
 
 from google.genai import types as genai_types
 from typing_extensions import override
@@ -30,6 +30,7 @@ class _EvalDatasetSchema(enum.Enum):
     """Represents the schema of an evaluation dataset."""
 
     GEMINI = "gemini"
+    FLATTEN = "flatten"
     UNKNOWN = "unknown"
 
 
@@ -141,3 +142,115 @@ def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset:
             eval_cases.append(eval_case)
 
         return types.EvaluationDataset(eval_cases=eval_cases)
+
+
+class _FlattenEvalDataConverter(_EvalDataConverter):
+    """Converter for datasets in a structured table format."""
+
+    def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset:
+        """Converts a list of raw data into an EvaluationDataset."""
+        eval_cases = []
+        for i, item_dict in enumerate(raw_data):
+            item = item_dict.copy()
+            eval_case_id = f"eval_case_{i}"
+            prompt_data = item.pop("prompt", None)
+            if not prompt_data:
+                prompt_data = item.pop("source", None)
+
+            conversation_history_data = item.pop("history", None)
+            response_data = item.pop("response", None)
+            reference_data = item.pop("reference", None)
+            system_instruction_data = item.pop("instruction", None)
+
+            if not response_data:
+                raise ValueError(
+                    f"Response is required but missing for {eval_case_id}."
+                )
+            if not prompt_data:
+                raise ValueError(f"Prompt is required but missing for {eval_case_id}.")
+
+            prompt: genai_types.Content
+            if isinstance(prompt_data, str):
+                prompt = genai_types.Content(parts=[genai_types.Part(text=prompt_data)])
+            elif isinstance(prompt_data, dict):
+                prompt = genai_types.Content.model_validate(prompt_data)
+            elif isinstance(prompt_data, genai_types.Content):
+                prompt = prompt_data
+            else:
+                raise ValueError(
+                    f"Invalid prompt type for case {i}: {type(prompt_data)}"
+                )
+
+            conversation_history: Optional[list[types.Message]] = None
+            if isinstance(conversation_history_data, list):
+                conversation_history = [
+                    types.Message(
+                        turn_id=str(turn_id),
+                        content=genai_types.Content.model_validate(content),
+                    )
+                    for turn_id, content in enumerate(conversation_history_data)
+                ]
+
+            responses: list[types.ResponseCandidate]
+            if isinstance(response_data, dict):
+                responses = [
+                    types.ResponseCandidate(
+                        response=genai_types.Content.model_validate(response_data)
+                    )
+                ]
+            elif isinstance(response_data, str):
+                responses = [
+                    types.ResponseCandidate(
+                        response=genai_types.Content(
+                            parts=[genai_types.Part(text=response_data)]
+                        )
+                    )
+                ]
+            elif isinstance(response_data, genai_types.Content):
+                responses = [types.ResponseCandidate(response=response_data)]
+            else:
+                raise ValueError(
+                    f"Invalid response type for case {i}: {type(response_data)}"
+                )
+
+            reference: Optional[types.ResponseCandidate] = None
+            if reference_data:
+                if isinstance(reference_data, dict):
+                    reference = types.ResponseCandidate(
+                        response=genai_types.Content.model_validate(reference_data)
+                    )
+                elif isinstance(reference_data, str):
+                    reference = types.ResponseCandidate(
+                        response=genai_types.Content(
+                            parts=[genai_types.Part(text=reference_data)]
+                        )
+                    )
+                elif isinstance(reference_data, genai_types.Content):
+                    reference = types.ResponseCandidate(response=reference_data)
+
+            system_instruction: Optional[genai_types.Content] = None
+            if system_instruction_data:
+                if isinstance(system_instruction_data, dict):
+                    system_instruction = genai_types.Content.model_validate(
+                        system_instruction_data
+                    )
+                elif isinstance(system_instruction_data, str):
+                    system_instruction = genai_types.Content(
+                        parts=[genai_types.Part(text=system_instruction_data)]
+                    )
+                elif isinstance(system_instruction_data, genai_types.Content):
+                    system_instruction = system_instruction_data
+
+            eval_case = types.EvalCase(
+                eval_case_id=eval_case_id,
+                prompt=prompt,
+                responses=responses,
+                reference=reference,
+                conversation_history=conversation_history,
+                system_instruction=system_instruction,
+                **item,  # Pass remaining columns as extra fields to EvalCase.
+                # They can be used for custom metric prompt templates.
+            )
+            eval_cases.append(eval_case)
+
+        return types.EvaluationDataset(eval_cases=eval_cases)

vertexai/_genai/types.py

@@ -23,7 +23,7 @@
 from typing import Any, Callable, ClassVar, Optional, Tuple, Union
 from google.genai import _common
 from google.genai import types as genai_types
-from pydantic import Field, computed_field, field_validator, model_validator
+from pydantic import ConfigDict, Field, computed_field, field_validator, model_validator
 from typing_extensions import TypedDict
 
 logger = logging.getLogger("vertexai_genai.types")
@@ -1909,6 +1909,8 @@ class EvalCase(_common.BaseModel):
         default=None,
         description="""Unique identifier for the evaluation case.""",
     )
+    # Allow extra fields to support custom metric prompts and stay backward compatible.
+    model_config = ConfigDict(frozen=True, extra="allow")
 
 
 class EvalCaseDict(TypedDict, total=False):