chore: Add MetricPromptBuilder type

PiperOrigin-RevId: 768286605

Author: jsondai

tests/unit/vertexai/genai/test_evals.py

@@ -703,6 +703,147 @@ def test_inference_with_multimodal_content(
         )
 
 
+class TestMetricPromptBuilder:
+    """Unit tests for the MetricPromptBuilder class."""
+
+    def test_metric_prompt_builder_minimal_fields(self):
+        criteria = {"criterion1": "definition1"}
+        rating_scores = {"score1": "description1"}
+        builder = vertexai_genai_types.MetricPromptBuilder(
+            criteria=criteria,
+            rating_scores=rating_scores,
+            metric_definition=None,
+            few_shot_examples=None,
+        )
+        assert builder.criteria == criteria
+        assert builder.rating_scores == rating_scores
+        expected_instruction = (
+            "You are an expert evaluator. Your task is to evaluate the quality"
+            " of the responses generated by AI models. We will provide you with"
+            " the user prompt and an AI-generated responses.\nYou should first"
+            " read the user input carefully for analyzing the task, and then"
+            " evaluate the quality of the responses based on the Criteria"
+            " provided in the Evaluation section below.\nYou will assign the"
+            " response a rating following the Rating Scores and Evaluation"
+            " Steps. Give step by step explanations for your rating, and only"
+            " choose ratings from the Rating Scores."
+        )
+        expected_evaluation_steps = {
+            "Step 1": (
+                "Assess the response in aspects of all criteria provided. Provide"
+                " assessment according to each criterion."
+            ),
+            "Step 2": (
+                "Score based on the Rating Scores. Give a brief rationale to"
+                " explain your evaluation considering each individual criterion."
+            ),
+        }
+        assert builder.instruction == expected_instruction
+        assert builder.evaluation_steps == expected_evaluation_steps
+        assert "{response}" in builder.text
+
+    def test_metric_prompt_builder_all_fields(self):
+        criteria = {"criterion1": "definition1"}
+        rating_scores = {"score1": "description1"}
+        instruction = "Custom instruction."
+        metric_definition = "Custom metric definition."
+        evaluation_steps = {"step1": "custom step 1"}
+        few_shot_examples = ["example1", "example2"]
+        builder = vertexai_genai_types.MetricPromptBuilder(
+            criteria=criteria,
+            rating_scores=rating_scores,
+            instruction=instruction,
+            metric_definition=metric_definition,
+            evaluation_steps=evaluation_steps,
+            few_shot_examples=few_shot_examples,
+        )
+        assert builder.criteria == criteria
+        assert builder.rating_scores == rating_scores
+        assert builder.instruction == instruction
+        assert builder.metric_definition == metric_definition
+        assert builder.evaluation_steps == evaluation_steps
+        assert builder.few_shot_examples == few_shot_examples
+        assert instruction in builder.text
+        assert metric_definition in builder.text
+        assert "custom step 1" in builder.text
+        assert "example1" in builder.text
+
+    def test_metric_prompt_builder_default_instruction_and_steps_in_text(self):
+        criteria = {"c1": "v1"}
+        rating_scores = {"s1": "d1"}
+        builder = vertexai_genai_types.MetricPromptBuilder(
+            criteria=criteria,
+            rating_scores=rating_scores,
+            metric_definition=None,
+            few_shot_examples=None,
+        )
+        expected_instruction = (
+            "You are an expert evaluator. Your task is to evaluate the quality"
+            " of the responses generated by AI models. We will provide you with"
+            " the user prompt and an AI-generated responses.\nYou should first"
+            " read the user input carefully for analyzing the task, and then"
+            " evaluate the quality of the responses based on the Criteria"
+            " provided in the Evaluation section below.\nYou will assign the"
+            " response a rating following the Rating Scores and Evaluation"
+            " Steps. Give step by step explanations for your rating, and only"
+            " choose ratings from the Rating Scores."
+        )
+        expected_evaluation_steps = {
+            "Step 1": (
+                "Assess the response in aspects of all criteria provided. Provide"
+                " assessment according to each criterion."
+            ),
+            "Step 2": (
+                "Score based on the Rating Scores. Give a brief rationale to"
+                " explain your evaluation considering each individual criterion."
+            ),
+        }
+        assert expected_instruction in builder.text
+        for step_desc in expected_evaluation_steps.values():
+            assert step_desc in builder.text
+
+    def test_metric_prompt_builder_custom_instruction_and_steps_in_text(self):
+        criteria = {"c1": "v1"}
+        rating_scores = {"s1": "d1"}
+        custom_instruction = "My custom instructions."
+        custom_steps = {"Step 1": "Do this first.", "Step 2": "Then do that."}
+        builder = vertexai_genai_types.MetricPromptBuilder(
+            criteria=criteria,
+            rating_scores=rating_scores,
+            instruction=custom_instruction,
+            evaluation_steps=custom_steps,
+            metric_definition=None,
+            few_shot_examples=None,
+        )
+        assert custom_instruction in builder.text
+        for step_desc in custom_steps.values():
+            assert step_desc in builder.text
+
+    def test_metric_prompt_builder_missing_criteria_raises_error(self):
+        with pytest.raises(
+            ValueError,
+            match="Both 'criteria' and 'rating_scores' are required",
+        ):
+            vertexai_genai_types.MetricPromptBuilder(
+                rating_scores={"score1": "description1"},
+                criteria=None,
+                metric_definition=None,
+                few_shot_examples=None,
+            )
+
+    def test_metric_prompt_builder_missing_rating_scores_raises_error(self):
+        with pytest.raises(
+            ValueError,
+            match="Both 'criteria' and 'rating_scores' are required",
+        ):
+            vertexai_genai_types.MetricPromptBuilder(
+                criteria={"criterion1": "definition1"},
+                rating_scores=None,
+                metric_definition=None,
+                few_shot_examples=None,
+            )
+
+
 class TestPromptTemplate:
     """Unit tests for the PromptTemplate class."""
 

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
+from pydantic import Field, computed_field, field_validator, model_validator
 from typing_extensions import TypedDict
 
 logger = logging.getLogger("vertexai_genai.types")
@@ -2463,6 +2463,173 @@ def __repr__(self) -> str:
         return f"PromptTemplate(text='{self.text}')"
 
 
+class MetricPromptBuilder(PromptTemplate):
+    """Builder class for structured LLM-based metric prompt template."""
+
+    criteria: Optional[dict[str, str]] = Field(
+        None,
+        description="""A dictionary of criteria used to evaluate the model responses.
+      The keys are criterion names, and the values are the corresponding
+      criterion definitions.
+      """,
+    )
+
+    rating_scores: Optional[dict[str, str]] = Field(
+        None,
+        description="""A dictionary mapping of rating score names to their definitions.""",
+    )
+
+    @staticmethod
+    def _get_default_instruction() -> str:
+        """Returns the default instruction for evaluation."""
+        return (
+            "You are an expert evaluator. Your task is to evaluate the quality"
+            " of the responses generated by AI models. We will provide you with"
+            " the user prompt and an AI-generated responses.\nYou should first"
+            " read the user input carefully for analyzing the task, and then"
+            " evaluate the quality of the responses based on the Criteria"
+            " provided in the Evaluation section below.\nYou will assign the"
+            " response a rating following the Rating Scores and Evaluation"
+            " Steps. Give step by step explanations for your rating, and only"
+            " choose ratings from the Rating Scores."
+        )
+
+    instruction: Optional[str] = Field(
+        default_factory=lambda: MetricPromptBuilder._get_default_instruction(),
+        description="""The general instruction to guide the model in performing the evaluation.
+    If not provided, a default instruction for evaluation will be used.
+    """,
+    )
+
+    metric_definition: Optional[str] = Field(
+        None,
+        description="""An optional high-level description of the metric to be evaluated.
+      If not provided, this field will not be included in the prompt template.
+      """,
+    )
+
+    @staticmethod
+    def _get_default_evaluation_steps() -> dict[str, str]:
+        """Returns the default evaluation steps for metric evaluation."""
+        return {
+            "Step 1": (
+                "Assess the response in aspects of all criteria provided."
+                " Provide assessment according to each criterion."
+            ),
+            "Step 2": (
+                "Score based on the Rating Scores. Give a brief rationale to"
+                " explain your evaluation considering each individual"
+                " criterion."
+            ),
+        }
+
+    evaluation_steps: Optional[dict[str, str]] = Field(
+        default_factory=lambda: MetricPromptBuilder._get_default_evaluation_steps(),
+        description="""An optional dictionary of evaluation steps.
+      The keys are the names of the evaluation steps, and the values are
+      descriptions of the corresponding evaluation steps. If not provided,
+      default metric evaluation steps will be used.
+      """,
+    )
+
+    few_shot_examples: Optional[list[str]] = Field(
+        None,
+        description="""An optional list of few-shot examples to guide the model's evaluation.
+      These examples demonstrate how to apply the criteria, rating scores,
+      and evaluation steps to assess model responses. Providing few-shot examples
+      can improve the accuracy of the evaluation. If not provided, this field
+      will not be included in the prompt template.
+      """,
+    )
+
+    @staticmethod
+    def _serialize_dict_in_order(elements: Optional[dict[str, str]]) -> str:
+        """Serializes dictionary to ordered string value without brackets."""
+        if elements is None:
+            return ""
+        return "\n".join(f"{key}: {value}" for key, value in sorted(elements.items()))
+
+    @model_validator(mode="before")
+    @classmethod
+    def _prepare_fields_and_construct_text(cls, data: Any) -> Any:
+        """Pydantic model validator (before mode) to prepare and construct prompt text.
+
+        This validator performs the following actions:
+          1. Apply default logic for fields (instruction, evaluation_steps).
+          2. Construct the 'text' string from all components.
+          3. Ensure 'text' is in the data dictionary for PromptTemplate
+          initialization.
+
+        Args:
+            data: Input data for the model, either a dictionary or an existing
+              model instance.
+
+        Returns:
+            Processed data dictionary with the 'text' field constructed.
+        """
+        if not isinstance(data, dict):
+            return data
+
+        if "text" in data:
+            raise ValueError(
+                "The 'text' field is automatically constructed and should not"
+                " be provided manually."
+            )
+
+        if data.get("criteria") is None or data.get("rating_scores") is None:
+            raise ValueError(
+                "Both 'criteria' and 'rating_scores' are required to construct"
+                " theLLM-based metric prompt template text."
+            )
+
+        instruction = data.get("instruction", cls._get_default_instruction())
+        metric_definition = data.get("metric_definition")
+        evaluation_steps = data.get(
+            "evaluation_steps", cls._get_default_evaluation_steps()
+        )
+        criteria = data.get("criteria")
+        rating_scores = data.get("rating_scores")
+        few_shot_examples = data.get("few_shot_examples")
+
+        template_parts = [
+            "# Instruction",
+            instruction,
+            "\n",
+            "# Evaluation",
+        ]
+
+        sections = {
+            "Metric Definition": metric_definition,
+            "Criteria": cls._serialize_dict_in_order(criteria),
+            "Rating Scores": cls._serialize_dict_in_order(rating_scores),
+            "Evaluation Steps": cls._serialize_dict_in_order(evaluation_steps),
+            "Evaluation Examples": (
+                "\n".join(few_shot_examples) if few_shot_examples else None
+            ),
+        }
+
+        for title, content in sections.items():
+            if content:
+                template_parts.extend([f"## {title}", f"{content}\n"])
+
+        template_parts.extend(
+            [
+                "\n# User Inputs and AI-generated Response",
+                "## User Inputs",
+            ]
+        )
+
+        template_parts.extend(["## AI-generated Response", "{response}"])
+        constructed_text = "\n".join(template_parts)
+
+        data["text"] = constructed_text
+        return data
+
+    def __str__(self) -> str:
+        """Returns the fully constructed prompt template text."""
+        return self.text
+
+
 class PromptTemplateDict(TypedDict, total=False):
     """A prompt template for creating prompts with variables."""