diff --git a/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py b/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py index d3085bec74..8db368c088 100644 --- a/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py +++ b/tests/unit/vertexai/genai/replays/test_public_generate_rubrics.py @@ -143,19 +143,21 @@ User prompt: {prompt}""" +_PROMPTS_DF = pd.DataFrame( + { + "prompt": [ + "Explain the theory of relativity in one sentence.", + "Write a short poem about a cat.", + ] + } +) + def test_public_method_generate_rubrics(client): """Tests the public generate_rubrics method.""" - prompts_df = pd.DataFrame( - { - "prompt": [ - "Explain the theory of relativity in one sentence.", - "Write a short poem about a cat.", - ] - } - ) + eval_dataset = client.evals.generate_rubrics( - src=prompts_df, + src=_PROMPTS_DF, prompt_template=_TEST_RUBRIC_GENERATION_PROMPT, rubric_group_name="text_quality_rubrics", ) @@ -176,6 +178,36 @@ def test_public_method_generate_rubrics(client): assert isinstance(first_rubric_group["text_quality_rubrics"][0], types.evals.Rubric) +def test_public_method_generate_rubrics_with_metric(client): + """Tests the public generate_rubrics method with a metric.""" + client._api_client._http_options.api_version = "v1beta1" + client._api_client._http_options.base_url = ( + "https://us-central1-staging-aiplatform.sandbox.googleapis.com/" + ) + metric_resource_name = "projects/977012026409/locations/us-central1/evaluationMetrics/6048334299558576128" + metric = types.Metric( + name="my_custom_metric", metric_resource_name=metric_resource_name + ) + eval_dataset = client.evals.generate_rubrics( + src=_PROMPTS_DF, rubric_group_name="my_registered_rubrics", metric=metric + ) + eval_dataset_df = eval_dataset.eval_dataset_df + + assert isinstance(eval_dataset, types.EvaluationDataset) + assert isinstance(eval_dataset_df, pd.DataFrame) + assert "rubric_groups" in eval_dataset_df.columns + assert len(eval_dataset_df) == 2 + + first_rubric_group = eval_dataset_df["rubric_groups"][0] + assert isinstance(first_rubric_group, dict) + assert "my_registered_rubrics" in first_rubric_group + assert isinstance(first_rubric_group["my_registered_rubrics"], list) + assert first_rubric_group["my_registered_rubrics"] + assert isinstance( + first_rubric_group["my_registered_rubrics"][0], types.evals.Rubric + ) + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py index 842db3412f..1b433d2f0f 100644 --- a/vertexai/_genai/evals.py +++ b/vertexai/_genai/evals.py @@ -526,6 +526,13 @@ def _GenerateInstanceRubricsRequest_to_vertex( ), ) + if getv(from_object, ["metric_resource_name"]) is not None: + setv( + to_object, + ["metricResourceName"], + getv(from_object, ["metric_resource_name"]), + ) + if getv(from_object, ["config"]) is not None: setv(to_object, ["config"], getv(from_object, ["config"])) @@ -1063,6 +1070,7 @@ def _generate_rubrics( types.PredefinedMetricSpecOrDict ] = None, rubric_generation_spec: Optional[types.RubricGenerationSpecOrDict] = None, + metric_resource_name: Optional[str] = None, config: Optional[types.RubricGenerationConfigOrDict] = None, ) -> types.GenerateInstanceRubricsResponse: """ @@ -1073,6 +1081,7 @@ def _generate_rubrics( contents=contents, predefined_rubric_generation_spec=predefined_rubric_generation_spec, rubric_generation_spec=rubric_generation_spec, + metric_resource_name=metric_resource_name, config=config, ) @@ -1575,16 +1584,20 @@ def generate_rubrics( rubric_type_ontology: Optional[list[str]] = None, predefined_spec_name: Optional[Union[str, "types.PrebuiltMetric"]] = None, metric_spec_parameters: Optional[dict[str, Any]] = None, + metric: Optional[types.MetricOrDict] = None, config: Optional[types.RubricGenerationConfigOrDict] = None, ) -> types.EvaluationDataset: """Generates rubrics for each prompt in the source and adds them as a new column structured as a dictionary. You can generate rubrics by providing either: - 1. A `predefined_spec_name` to use a Vertex AI backend recipe. - 2. A `prompt_template` along with other configuration parameters + 1. A `metric` to use a pre-registered metric resource. + 2. A `predefined_spec_name` to use a Vertex AI backend recipe. + 3. A `prompt_template` along with other configuration parameters (`generator_model_config`, `rubric_content_type`, `rubric_type_ontology`) for custom rubric generation. + with `metric` taking precedence over `predefined_spec_name`, + and `predefined_spec_name` taking precedence over `prompt_template` These two modes are mutually exclusive. @@ -1614,6 +1627,9 @@ def generate_rubrics( metric_spec_parameters: Optional. Parameters for the Predefined Metric, used to customize rubric generation. Only used if `predefined_spec_name` is set. Example: {"guidelines": ["The response must be in Japanese."]} + metric: Optional. A types.Metric object containing a metric_resource_name, + or a resource name string. If provided, this will take precedence over + predefined_spec_name and prompt_template. config: Optional. Configuration for the rubric generation process. Returns: @@ -1653,10 +1669,32 @@ def generate_rubrics( ) all_rubric_groups: list[dict[str, list[types.Rubric]]] = [] + actual_metric_resource_name = None + if metric: + if isinstance(metric, str) and metric.startswith("projects/"): + actual_metric_resource_name = metric + else: + metric_obj = ( + types.Metric.model_validate(metric) + if isinstance(metric, dict) + else metric + ) + actual_metric_resource_name = getattr( + metric_obj, "metric_resource_name", None + ) + if not actual_metric_resource_name: + raise ValueError( + "The provided Metric object must have metric_resource_name set." + ) + rubric_gen_spec = None predefined_spec = None - if predefined_spec_name: + if actual_metric_resource_name: + # Precedence: Registered metric resource overrides everything else. + predefined_spec = None + rubric_gen_spec = None + elif predefined_spec_name: if prompt_template: logger.warning( "prompt_template is ignored when predefined_spec_name is provided." @@ -1713,7 +1751,7 @@ def generate_rubrics( rubric_gen_spec = types.RubricGenerationSpec.model_validate(spec_dict) else: raise ValueError( - "Either predefined_spec_name or prompt_template must be provided." + "Either metric, predefined_spec_name or prompt_template must be provided." ) for _, row in prompts_df.iterrows(): @@ -1736,6 +1774,7 @@ def generate_rubrics( contents=contents, rubric_generation_spec=rubric_gen_spec, predefined_rubric_generation_spec=predefined_spec, + metric_resource_name=actual_metric_resource_name, config=config, ) rubric_group = {rubric_group_name: response.generated_rubrics} @@ -2321,6 +2360,7 @@ async def _generate_rubrics( types.PredefinedMetricSpecOrDict ] = None, rubric_generation_spec: Optional[types.RubricGenerationSpecOrDict] = None, + metric_resource_name: Optional[str] = None, config: Optional[types.RubricGenerationConfigOrDict] = None, ) -> types.GenerateInstanceRubricsResponse: """ @@ -2331,6 +2371,7 @@ async def _generate_rubrics( contents=contents, predefined_rubric_generation_spec=predefined_rubric_generation_spec, rubric_generation_spec=rubric_generation_spec, + metric_resource_name=metric_resource_name, config=config, ) diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index ef1549ea75..a33f015284 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -5368,6 +5368,10 @@ class _GenerateInstanceRubricsRequest(_common.BaseModel): default=None, description="""Specification for how the rubrics should be generated.""", ) + metric_resource_name: Optional[str] = Field( + default=None, + description="""Registered metric resource name. If this field is set, the configuration provided in this field is used for rubric generation. The `predefined_rubric_generation_spec` and `rubric_generation_spec` fields will be ignored.""", + ) config: Optional[RubricGenerationConfig] = Field(default=None, description="""""") @@ -5388,6 +5392,9 @@ class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): rubric_generation_spec: Optional[RubricGenerationSpecDict] """Specification for how the rubrics should be generated.""" + metric_resource_name: Optional[str] + """Registered metric resource name. If this field is set, the configuration provided in this field is used for rubric generation. The `predefined_rubric_generation_spec` and `rubric_generation_spec` fields will be ignored.""" + config: Optional[RubricGenerationConfigDict] """"""