fix: GenAI Eval: fix a prompt template variable parsing issue when multiline json structures are present

PiperOrigin-RevId: 738555371

Author: jsondai

tests/unit/vertexai/test_evaluation.py

@@ -2041,6 +2041,61 @@ def test_initialize_metric_column_mapping(self):
         )
         assert converted_metric_column_mapping == _EXPECTED_COLUMN_MAPPING
 
+    def test_upload_results(self, mock_storage_blob_from_string):
+        with mock.patch("json.dump") as mock_json_dump:
+            evaluation.utils.upload_evaluation_results(
+                MOCK_EVAL_RESULT,
+                _TEST_BUCKET,
+                _TEST_FILE_NAME,
+                "candidate_model",
+                "baseline_model",
+                "gs://test-bucket/test-dataset.csv",
+                [_TEST_POINTWISE_METRIC, _TEST_PAIRWISE_METRIC],
+            )
+
+        mock_storage_blob_from_string.assert_any_call(
+            uri="gs://test-bucket/test-file-name/test-file-name.csv",
+            client=mock.ANY,
+        )
+        mock_storage_blob_from_string.assert_any_call(
+            uri="gs://test-bucket/test-file-name/summary_metrics.json",
+            client=mock.ANY,
+        )
+        mock_json_dump.assert_called_once_with(
+            {
+                "summary_metrics": MOCK_EVAL_RESULT.summary_metrics,
+                "candidate_model_name": "candidate_model",
+                "baseline_model_name": "baseline_model",
+                "dataset_uri": "gs://test-bucket/test-dataset.csv",
+                "metric_descriptions": {
+                    "test_pointwise_metric": {
+                        "criteria": _CRITERIA,
+                        "rating_rubric": _POINTWISE_RATING_RUBRIC,
+                    },
+                    "test_pairwise_metric": {
+                        "criteria": _CRITERIA,
+                        "rating_rubric": _PAIRWISE_RATING_RUBRIC,
+                    },
+                },
+            },
+            mock.ANY,
+        )
+
+    def test_upload_results_with_default_file_name(self, mock_storage_blob_from_string):
+        with mock.patch.object(
+            aiplatform_utils, "timestamped_unique_name"
+        ) as mock_timestamped_unique_name:
+            mock_timestamped_unique_name.return_value = "2025-02-10-12-00-00-12345"
+            evaluation.utils.upload_evaluation_results(
+                MOCK_EVAL_RESULT,
+                _TEST_BUCKET,
+            )
+
+        mock_storage_blob_from_string.assert_any_call(
+            uri="gs://test-bucket/eval_results_2025-02-10-12-00-00-12345/eval_results_2025-02-10-12-00-00-12345.csv",
+            client=mock.ANY,
+        )
+
 
 class TestPromptTemplate:
     def test_init(self):
@@ -2138,57 +2193,31 @@ def test_pairtwise_metric_prompt_template_with_default_values(self):
             == _EXPECTED_PAIRWISE_PROMPT_TEMPLATE_WITH_DEFAULT_VALUES.strip()
         )
 
-    def test_upload_results(self, mock_storage_blob_from_string):
-        with mock.patch("json.dump") as mock_json_dump:
-            evaluation.utils.upload_evaluation_results(
-                MOCK_EVAL_RESULT,
-                _TEST_BUCKET,
-                _TEST_FILE_NAME,
-                "candidate_model",
-                "baseline_model",
-                "gs://test-bucket/test-dataset.csv",
-                [_TEST_POINTWISE_METRIC, _TEST_PAIRWISE_METRIC],
-            )
-
-        mock_storage_blob_from_string.assert_any_call(
-            uri="gs://test-bucket/test-file-name/test-file-name.csv",
-            client=mock.ANY,
-        )
-        mock_storage_blob_from_string.assert_any_call(
-            uri="gs://test-bucket/test-file-name/summary_metrics.json",
-            client=mock.ANY,
-        )
-        mock_json_dump.assert_called_once_with(
-            {
-                "summary_metrics": MOCK_EVAL_RESULT.summary_metrics,
-                "candidate_model_name": "candidate_model",
-                "baseline_model_name": "baseline_model",
-                "dataset_uri": "gs://test-bucket/test-dataset.csv",
-                "metric_descriptions": {
-                    "test_pointwise_metric": {
-                        "criteria": _CRITERIA,
-                        "rating_rubric": _POINTWISE_RATING_RUBRIC,
-                    },
-                    "test_pairwise_metric": {
-                        "criteria": _CRITERIA,
-                        "rating_rubric": _PAIRWISE_RATING_RUBRIC,
-                    },
-                },
-            },
-            mock.ANY,
-        )
-
-    def test_upload_results_with_default_file_name(self, mock_storage_blob_from_string):
-        with mock.patch.object(
-            aiplatform_utils, "timestamped_unique_name"
-        ) as mock_timestamped_unique_name:
-            mock_timestamped_unique_name.return_value = "2025-02-10-12-00-00-12345"
-            evaluation.utils.upload_evaluation_results(
-                MOCK_EVAL_RESULT,
-                _TEST_BUCKET,
-            )
-
-        mock_storage_blob_from_string.assert_any_call(
-            uri="gs://test-bucket/eval_results_2025-02-10-12-00-00-12345/eval_results_2025-02-10-12-00-00-12345.csv",
-            client=mock.ANY,
-        )
+    def test_complex_prompt_template_variables(self):
+        template_str = """Metric prompt template
+instructions ...
+Here are some JSON structures
+{
+  "Function API spec": You may use default python libraries,
+  "example": test test
+}
+Output format prompt with JSON:
+The answer should be a json alone which follows the json structure below:
+{
+  "is_the_response_valid": [valid or invalid],
+  "reasoning":
+  "rewritten response":
+}
+Here are some actual variables:
+{var_1} {var2} {_var_3} {
+        var_5_mutli_line
+} {VAR_6} {7_var} {{var_9}}
+"""
+        prompt_template = evaluation.PromptTemplate(template_str)
+        assert prompt_template.variables == {
+            "var_1",
+            "var2",
+            "_var_3",
+            "VAR_6",
+            "var_9",
+        }

vertexai/evaluation/prompt_template.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,18 +14,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import string
+"""Prompt template for creating prompts with variables."""
+
+import re
 from typing import Set
 
+_VARIABLE_NAME_REGEX = r"\{([_a-zA-Z][_a-zA-Z0-9]*)\}"
+
 
 class PromptTemplate:
     """A prompt template for creating prompts with variables.
 
     The `PromptTemplate` class allows users to define a template string with
     variables represented in curly braces `{variable}`. The variable
-    names cannot contain spaces. These variables can be replaced with specific
-    values using the `assemble` method, providing flexibility in generating
-    dynamic prompts.
+    names cannot contain spaces and must start with a letter or underscore,
+    followed by letters, digits, or underscore. These variables can be
+    replaced with specific values using the `assemble` method, providing
+    flexibility in generating dynamic prompts.
 
     Usage:
 
@@ -49,11 +54,7 @@ def __init__(self, template: str):
 
     def _get_variables(self) -> Set[str]:
         """Extracts and return a set of variable names from the template."""
-        return set(
-            field_name
-            for _, field_name, _, _ in string.Formatter().parse(self.template)
-            if field_name is not None
-        )
+        return set(re.findall(_VARIABLE_NAME_REGEX, self.template))
 
     def assemble(self, **kwargs) -> "PromptTemplate":
         """Replaces only the provided variables in the template with specific values.

vertexai/preview/evaluation/prompt_template.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Copyright 2024 Google LLC
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,18 +16,21 @@
 #
 """Prompt template for creating prompts with variables."""
 
-import string
+import re
 from typing import Set
 
+_VARIABLE_NAME_REGEX = r"\{([_a-zA-Z][_a-zA-Z0-9]*)\}"
+
 
 class PromptTemplate:
     """A prompt template for creating prompts with variables.
 
     The `PromptTemplate` class allows users to define a template string with
     variables represented in curly braces `{variable}`. The variable
-    names cannot contain spaces. These variables can be replaced with specific
-    values using the `assemble` method, providing flexibility in generating
-    dynamic prompts.
+    names cannot contain spaces and must start with a letter or underscore,
+    followed by letters, digits, or underscore. These variables can be
+    replaced with specific values using the `assemble` method, providing
+    flexibility in generating dynamic prompts.
 
     Usage:
 
@@ -51,11 +54,7 @@ def __init__(self, template: str):
 
     def _get_variables(self) -> Set[str]:
         """Extracts and return a set of variable names from the template."""
-        return set(
-            field_name
-            for _, field_name, _, _ in string.Formatter().parse(self.template)
-            if field_name is not None
-        )
+        return set(re.findall(_VARIABLE_NAME_REGEX, self.template))
 
     def assemble(self, **kwargs) -> "PromptTemplate":
         """Replaces only the provided variables in the template with specific values.