From b7a861404301008729d1dee8698700f003ffc888 Mon Sep 17 00:00:00 2001
From: Christian Yarros <yarros@google.com>
Date: Sun, 1 Sep 2024 22:09:52 +0000
Subject: [PATCH] Add RunEvaluationOperator for Google Vertex AI Rapid
 Evaluation API

---
 .../cloud/hooks/vertex_ai/generative_model.py | 90 ++++++++++++++++-
 .../operators/vertex_ai/generative_model.py   | 99 +++++++++++++++++++
 airflow/providers/google/provider.yaml        |  1 +
 .../operators/cloud/vertex_ai.rst             | 10 ++
 generated/provider_dependencies.json          |  1 +
 .../hooks/vertex_ai/test_generative_model.py  | 67 ++++++++++++-
 .../vertex_ai/test_generative_model.py        | 89 +++++++++++++++++
 .../example_vertex_ai_generative_model.py     | 47 +++++++++
 8 files changed, 400 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py b/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py
index 2fe73221bfd53..acadbb788a905 100644
--- a/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py
+++ b/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py
@@ -25,6 +25,7 @@
 import vertexai
 from vertexai.generative_models import GenerativeModel, Part
 from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
+from vertexai.preview.evaluation import EvalResult, EvalTask
 from vertexai.preview.tuning import sft
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -62,11 +63,38 @@ def get_text_embedding_model(self, pretrained_model: str):
         model = TextEmbeddingModel.from_pretrained(pretrained_model)
         return model
 
-    def get_generative_model(self, pretrained_model: str) -> GenerativeModel:
+    def get_generative_model(
+        self,
+        pretrained_model: str,
+        system_instruction: str | None = None,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
+        tools: list | None = None,
+    ) -> GenerativeModel:
         """Return a Generative Model object."""
-        model = GenerativeModel(pretrained_model)
+        model = GenerativeModel(
+            model_name=pretrained_model,
+            system_instruction=system_instruction,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+            tools=tools,
+        )
         return model
 
+    def get_eval_task(
+        self,
+        dataset: dict,
+        metrics: list,
+        experiment: str,
+    ) -> EvalTask:
+        """Return an EvalTask object."""
+        eval_task = EvalTask(
+            dataset=dataset,
+            metrics=metrics,
+            experiment=experiment,
+        )
+        return eval_task
+
     @deprecated(
         planned_removal_date="January 01, 2025",
         use_instead="Part objects included in contents parameter of "
@@ -436,3 +464,61 @@ def count_tokens(
         )
 
         return response
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def run_evaluation(
+        self,
+        pretrained_model: str,
+        eval_dataset: dict,
+        metrics: list,
+        experiment_name: str,
+        experiment_run_name: str,
+        prompt_template: str,
+        location: str,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
+        system_instruction: str | None = None,
+        tools: list | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
+    ) -> EvalResult:
+        """
+        Use the Rapid Evaluation API to evaluate a model.
+
+        :param pretrained_model: Required. A pre-trained model optimized for performing natural
+            language tasks such as classification, summarization, extraction, content
+            creation, and ideation.
+        :param eval_dataset: Required. A fixed dataset for evaluating a model against. Adheres to Rapid Evaluation API.
+        :param metrics: Required. A list of evaluation metrics to be used in the experiment. Adheres to Rapid Evaluation API.
+        :param experiment_name: Required. The name of the evaluation experiment.
+        :param experiment_run_name: Required. The specific run name or ID for this experiment.
+        :param prompt_template: Required. The template used to format the model's prompts during evaluation. Adheres to Rapid Evaluation API.
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param location: Required. The ID of the Google Cloud location that the service belongs to.
+        :param generation_config: Optional. A dictionary containing generation parameters for the model.
+        :param safety_settings: Optional. A dictionary specifying harm category thresholds for blocking model outputs.
+        :param system_instruction: Optional. An instruction given to the model to guide its behavior.
+        :param tools: Optional. A list of tools available to the model during evaluation, such as a data store.
+        """
+        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
+
+        model = self.get_generative_model(
+            pretrained_model=pretrained_model,
+            system_instruction=system_instruction,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+            tools=tools,
+        )
+
+        eval_task = self.get_eval_task(
+            dataset=eval_dataset,
+            metrics=metrics,
+            experiment=experiment_name,
+        )
+
+        eval_result = eval_task.evaluate(
+            model=model,
+            prompt_template=prompt_template,
+            experiment_run_name=experiment_run_name,
+        )
+
+        return eval_result
diff --git a/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py b/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py
index 3a49035002a38..b0b9462e2151d 100644
--- a/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py
+++ b/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py
@@ -736,3 +736,102 @@ def execute(self, context: Context):
         self.xcom_push(context, key="total_billable_characters", value=response.total_billable_characters)
 
         return types_v1beta1.CountTokensResponse.to_dict(response)
+
+
+class RunEvaluationOperator(GoogleCloudBaseOperator):
+    """
+    Use the Rapid Evaluation API to evaluate a model.
+
+    :param pretrained_model: Required. A pre-trained model optimized for performing natural
+        language tasks such as classification, summarization, extraction, content
+        creation, and ideation.
+    :param eval_dataset: Required. A fixed dataset for evaluating a model against. Adheres to Rapid Evaluation API.
+    :param metrics: Required. A list of evaluation metrics to be used in the experiment. Adheres to Rapid Evaluation API.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param experiment_run_name: Required. The specific run name or ID for this experiment.
+    :param prompt_template: Required. The template used to format the model's prompts during evaluation. Adheres to Rapid Evaluation API.
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param generation_config: Optional. A dictionary containing generation parameters for the model.
+    :param safety_settings: Optional. A dictionary specifying harm category thresholds for blocking model outputs.
+    :param system_instruction: Optional. An instruction given to the model to guide its behavior.
+    :param tools: Optional. A list of tools available to the model during evaluation, such as a data store.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "pretrained_model",
+        "eval_dataset",
+        "prompt_template",
+        "experiment_name",
+        "experiment_run_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        pretrained_model: str,
+        eval_dataset: dict,
+        metrics: list,
+        experiment_name: str,
+        experiment_run_name: str,
+        prompt_template: str,
+        project_id: str,
+        location: str,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
+        system_instruction: str | None = None,
+        tools: list | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        self.pretrained_model = pretrained_model
+        self.eval_dataset = eval_dataset
+        self.metrics = metrics
+        self.experiment_name = experiment_name
+        self.experiment_run_name = experiment_run_name
+        self.prompt_template = prompt_template
+        self.system_instruction = system_instruction
+        self.generation_config = generation_config
+        self.safety_settings = safety_settings
+        self.tools = tools
+        self.project_id = project_id
+        self.location = location
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        response = self.hook.run_evaluation(
+            pretrained_model=self.pretrained_model,
+            eval_dataset=self.eval_dataset,
+            metrics=self.metrics,
+            experiment_name=self.experiment_name,
+            experiment_run_name=self.experiment_run_name,
+            prompt_template=self.prompt_template,
+            project_id=self.project_id,
+            location=self.location,
+            system_instruction=self.system_instruction,
+            generation_config=self.generation_config,
+            safety_settings=self.safety_settings,
+            tools=self.tools,
+        )
+
+        return response.summary_metrics
diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml
index 6311001448e65..46d82050ffc1c 100644
--- a/airflow/providers/google/provider.yaml
+++ b/airflow/providers/google/provider.yaml
@@ -169,6 +169,7 @@ dependencies:
   - sqlalchemy-bigquery>=1.2.1
   - sqlalchemy-spanner>=1.6.2
   - tenacity>=8.1.0
+  - immutabledict>=4.2.0
 
 additional-extras:
   - name: apache.beam
diff --git a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst
index 87c7b81bc96e9..db23d4dfeabb2 100644
--- a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst
+++ b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst
@@ -636,6 +636,16 @@ The operator returns the total tokens in :ref:`XCom <concepts:xcom>` under ``tot
     :start-after: [START how_to_cloud_vertex_ai_count_tokens_operator]
     :end-before: [END how_to_cloud_vertex_ai_count_tokens_operator]
 
+To evaluate a model you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.generative_model.RunEvaluationOperator`.
+The operator returns the evaluation summary metrics in :ref:`XCom <concepts:xcom>` under ``summary_metrics`` key.
+
+.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_run_evaluation_operator]
+    :end-before: [END how_to_cloud_vertex_ai_run_evaluation_operator]
+
 Reference
 ^^^^^^^^^
 
diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json
index afff3d8afbb9f..cc1f796f75fde 100644
--- a/generated/provider_dependencies.json
+++ b/generated/provider_dependencies.json
@@ -655,6 +655,7 @@
       "google-cloud-workflows>=1.10.0",
       "grpcio-gcp>=0.2.2",
       "httpx>=0.25.0",
+      "immutabledict>=4.2.0",
       "json-merge-patch>=0.2",
       "looker-sdk>=22.4.0",
       "pandas-gbq>=0.7.0",
diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py b/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py
index 5eff60b91d458..2004eec29fef1 100644
--- a/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py
+++ b/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py
@@ -27,8 +27,8 @@
 
 # For no Pydantic environment, we need to skip the tests
 pytest.importorskip("google.cloud.aiplatform_v1")
-vertexai = pytest.importorskip("vertexai.generative_models")
 from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding
+from vertexai.preview.evaluation import MetricPromptTemplateExamples
 
 from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import (
     GenerativeModelHook,
@@ -73,6 +73,38 @@
 SOURCE_MODEL = "gemini-1.0-pro-002"
 TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl"
 
+TEST_EVAL_DATASET = {
+    "context": [
+        "To make a classic spaghetti carbonara, start by bringing a large pot of salted water to a boil. While the water is heating up, cook pancetta or guanciale in a skillet with olive oil over medium heat until it's crispy and golden brown. Once the pancetta is done, remove it from the skillet and set it aside. In the same skillet, whisk together eggs, grated Parmesan cheese, and black pepper to make the sauce. When the pasta is cooked al dente, drain it and immediately toss it in the skillet with the egg mixture, adding a splash of the pasta cooking water to create a creamy sauce.",
+        "Preparing a perfect risotto requires patience and attention to detail. Begin by heating butter in a large, heavy-bottomed pot over medium heat. Add finely chopped onions and minced garlic to the pot, and cook until they're soft and translucent, about 5 minutes. Next, add Arborio rice to the pot and cook, stirring constantly, until the grains are coated with the butter and begin to toast slightly. Pour in a splash of white wine and cook until it's absorbed. From there, gradually add hot chicken or vegetable broth to the rice, stirring frequently, until the risotto is creamy and the rice is tender with a slight bite.",
+        "For a flavorful grilled steak, start by choosing a well-marbled cut of beef like ribeye or New York strip. Season the steak generously with kosher salt and freshly ground black pepper on both sides, pressing the seasoning into the meat. Preheat a grill to high heat and brush the grates with oil to prevent sticking. Place the seasoned steak on the grill and cook for about 4-5 minutes on each side for medium-rare, or adjust the cooking time to your desired level of doneness. Let the steak rest for a few minutes before slicing against the grain and serving.",
+        "Creating a creamy homemade tomato soup is a comforting and simple process. Begin by heating olive oil in a large pot over medium heat. Add diced onions and minced garlic to the pot and cook until they're soft and fragrant. Next, add chopped fresh tomatoes, chicken or vegetable broth, and a sprig of fresh basil to the pot. Simmer the soup for about 20-30 minutes, or until the tomatoes are tender and falling apart. Remove the basil sprig and use an immersion blender to puree the soup until smooth. Season with salt and pepper to taste before serving.",
+        "To bake a decadent chocolate cake from scratch, start by preheating your oven to 350°F (175°C) and greasing and flouring two 9-inch round cake pans. In a large mixing bowl, cream together softened butter and granulated sugar until light and fluffy. Beat in eggs one at a time, making sure each egg is fully incorporated before adding the next. In a separate bowl, sift together all-purpose flour, cocoa powder, baking powder, baking soda, and salt. Divide the batter evenly between the prepared cake pans and bake for 25-30 minutes, or until a toothpick inserted into the center comes out clean.",
+    ],
+    "instruction": ["Summarize the following article"] * 5,
+    "reference": [
+        "The process of making spaghetti carbonara involves boiling pasta, crisping pancetta or guanciale, whisking together eggs and Parmesan cheese, and tossing everything together to create a creamy sauce.",
+        "Preparing risotto entails sautéing onions and garlic, toasting Arborio rice, adding wine and broth gradually, and stirring until creamy and tender.",
+        "Grilling a flavorful steak involves seasoning generously, preheating the grill, cooking to desired doneness, and letting it rest before slicing.",
+        "Creating homemade tomato soup includes sautéing onions and garlic, simmering with tomatoes and broth, pureeing until smooth, and seasoning to taste.",
+        "Baking a decadent chocolate cake requires creaming butter and sugar, beating in eggs and alternating dry ingredients with buttermilk before baking until done.",
+    ],
+}
+TEST_METRICS = [
+    MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY,
+    MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS,
+    MetricPromptTemplateExamples.Pointwise.VERBOSITY,
+    MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING,
+    "exact_match",
+    "bleu",
+    "rouge_1",
+    "rouge_2",
+    "rouge_l_sum",
+]
+TEST_EXPERIMENT_NAME = "eval-experiment-airflow-operator"
+TEST_EXPERIMENT_RUN_NAME = "eval-experiment-airflow-operator-run"
+TEST_PROMPT_TEMPLATE = "{instruction}. Article: {context}. Summary:"
+
 BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}"
 GENERATIVE_MODEL_STRING = "airflow.providers.google.cloud.hooks.vertex_ai.generative_model.{}"
 
@@ -207,7 +239,6 @@ def test_supervised_fine_tuning_train(self, mock_sft_train) -> None:
             train_dataset=TRAIN_DATASET,
         )
 
-        # Assertions
         mock_sft_train.assert_called_once_with(
             source_model=SOURCE_MODEL,
             train_dataset=TRAIN_DATASET,
@@ -230,3 +261,35 @@ def test_count_tokens(self, mock_model) -> None:
         mock_model.return_value.count_tokens.assert_called_once_with(
             contents=TEST_CONTENTS,
         )
+
+    @mock.patch(GENERATIVE_MODEL_STRING.format("GenerativeModelHook.get_generative_model"))
+    @mock.patch(GENERATIVE_MODEL_STRING.format("GenerativeModelHook.get_eval_task"))
+    def test_run_evaluation(self, mock_eval_task, mock_model) -> None:
+        self.hook.run_evaluation(
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            pretrained_model=TEST_MULTIMODAL_PRETRAINED_MODEL,
+            eval_dataset=TEST_EVAL_DATASET,
+            metrics=TEST_METRICS,
+            experiment_name=TEST_EXPERIMENT_NAME,
+            experiment_run_name=TEST_EXPERIMENT_RUN_NAME,
+            prompt_template=TEST_PROMPT_TEMPLATE,
+        )
+
+        mock_model.assert_called_once_with(
+            pretrained_model=TEST_MULTIMODAL_PRETRAINED_MODEL,
+            system_instruction=None,
+            generation_config=None,
+            safety_settings=None,
+            tools=None,
+        )
+        mock_eval_task.assert_called_once_with(
+            dataset=TEST_EVAL_DATASET,
+            metrics=TEST_METRICS,
+            experiment=TEST_EXPERIMENT_NAME,
+        )
+        mock_eval_task.return_value.evaluate.assert_called_once_with(
+            model=mock_model.return_value,
+            prompt_template=TEST_PROMPT_TEMPLATE,
+            experiment_run_name=TEST_EXPERIMENT_RUN_NAME,
+        )
diff --git a/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py b/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py
index ada9d7843df93..3745850e72108 100644
--- a/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py
+++ b/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py
@@ -29,6 +29,7 @@
 pytest.importorskip("google.cloud.aiplatform_v1beta1")
 vertexai = pytest.importorskip("vertexai.generative_models")
 from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding
+from vertexai.preview.evaluation import MetricPromptTemplateExamples
 
 from airflow.providers.google.cloud.operators.vertex_ai.generative_model import (
     CountTokensOperator,
@@ -37,6 +38,7 @@
     PromptLanguageModelOperator,
     PromptMultimodalModelOperator,
     PromptMultimodalModelWithMediaOperator,
+    RunEvaluationOperator,
     SupervisedFineTuningTrainOperator,
     TextEmbeddingModelGetEmbeddingsOperator,
     TextGenerationModelPredictOperator,
@@ -448,3 +450,90 @@ def test_execute(self, to_dict_mock, mock_hook):
             contents=contents,
             pretrained_model=pretrained_model,
         )
+
+
+class TestVertexAIRunEvaluationOperator:
+    @mock.patch(VERTEX_AI_PATH.format("generative_model.GenerativeModelHook"))
+    def test_execute(
+        self,
+        mock_hook,
+    ):
+        tools = [Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval())]
+        pretrained_model = "gemini-pro"
+        safety_settings = {
+            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+        }
+        generation_config = {"max_output_tokens": 256, "top_p": 0.8, "temperature": 0.0}
+
+        eval_dataset = {
+            "context": [
+                "To make a classic spaghetti carbonara, start by bringing a large pot of salted water to a boil. While the water is heating up, cook pancetta or guanciale in a skillet with olive oil over medium heat until it's crispy and golden brown. Once the pancetta is done, remove it from the skillet and set it aside. In the same skillet, whisk together eggs, grated Parmesan cheese, and black pepper to make the sauce. When the pasta is cooked al dente, drain it and immediately toss it in the skillet with the egg mixture, adding a splash of the pasta cooking water to create a creamy sauce.",
+                "Preparing a perfect risotto requires patience and attention to detail. Begin by heating butter in a large, heavy-bottomed pot over medium heat. Add finely chopped onions and minced garlic to the pot, and cook until they're soft and translucent, about 5 minutes. Next, add Arborio rice to the pot and cook, stirring constantly, until the grains are coated with the butter and begin to toast slightly. Pour in a splash of white wine and cook until it's absorbed. From there, gradually add hot chicken or vegetable broth to the rice, stirring frequently, until the risotto is creamy and the rice is tender with a slight bite.",
+                "For a flavorful grilled steak, start by choosing a well-marbled cut of beef like ribeye or New York strip. Season the steak generously with kosher salt and freshly ground black pepper on both sides, pressing the seasoning into the meat. Preheat a grill to high heat and brush the grates with oil to prevent sticking. Place the seasoned steak on the grill and cook for about 4-5 minutes on each side for medium-rare, or adjust the cooking time to your desired level of doneness. Let the steak rest for a few minutes before slicing against the grain and serving.",
+                "Creating a creamy homemade tomato soup is a comforting and simple process. Begin by heating olive oil in a large pot over medium heat. Add diced onions and minced garlic to the pot and cook until they're soft and fragrant. Next, add chopped fresh tomatoes, chicken or vegetable broth, and a sprig of fresh basil to the pot. Simmer the soup for about 20-30 minutes, or until the tomatoes are tender and falling apart. Remove the basil sprig and use an immersion blender to puree the soup until smooth. Season with salt and pepper to taste before serving.",
+                "To bake a decadent chocolate cake from scratch, start by preheating your oven to 350°F (175°C) and greasing and flouring two 9-inch round cake pans. In a large mixing bowl, cream together softened butter and granulated sugar until light and fluffy. Beat in eggs one at a time, making sure each egg is fully incorporated before adding the next. In a separate bowl, sift together all-purpose flour, cocoa powder, baking powder, baking soda, and salt. Divide the batter evenly between the prepared cake pans and bake for 25-30 minutes, or until a toothpick inserted into the center comes out clean.",
+            ],
+            "instruction": ["Summarize the following article"] * 5,
+            "reference": [
+                "The process of making spaghetti carbonara involves boiling pasta, crisping pancetta or guanciale, whisking together eggs and Parmesan cheese, and tossing everything together to create a creamy sauce.",
+                "Preparing risotto entails sautéing onions and garlic, toasting Arborio rice, adding wine and broth gradually, and stirring until creamy and tender.",
+                "Grilling a flavorful steak involves seasoning generously, preheating the grill, cooking to desired doneness, and letting it rest before slicing.",
+                "Creating homemade tomato soup includes sautéing onions and garlic, simmering with tomatoes and broth, pureeing until smooth, and seasoning to taste.",
+                "Baking a decadent chocolate cake requires creaming butter and sugar, beating in eggs and alternating dry ingredients with buttermilk before baking until done.",
+            ],
+        }
+        metrics = [
+            MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY,
+            MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS,
+            MetricPromptTemplateExamples.Pointwise.VERBOSITY,
+            MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING,
+            "exact_match",
+            "bleu",
+            "rouge_1",
+            "rouge_2",
+            "rouge_l_sum",
+        ]
+        experiment_name = "eval-experiment-airflow-operator"
+        experiment_run_name = "eval-experiment-airflow-operator-run"
+        prompt_template = "{instruction}. Article: {context}. Summary:"
+        system_instruction = "be concise."
+
+        op = RunEvaluationOperator(
+            task_id=TASK_ID,
+            pretrained_model=pretrained_model,
+            eval_dataset=eval_dataset,
+            metrics=metrics,
+            experiment_name=experiment_name,
+            experiment_run_name=experiment_run_name,
+            prompt_template=prompt_template,
+            system_instruction=system_instruction,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+            tools=tools,
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+        op.execute(context={"ti": mock.MagicMock()})
+        mock_hook.assert_called_once_with(
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+        )
+        mock_hook.return_value.run_evaluation.assert_called_once_with(
+            project_id=GCP_PROJECT,
+            location=GCP_LOCATION,
+            pretrained_model=pretrained_model,
+            eval_dataset=eval_dataset,
+            metrics=metrics,
+            experiment_name=experiment_name,
+            experiment_run_name=experiment_run_name,
+            prompt_template=prompt_template,
+            system_instruction=system_instruction,
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+            tools=tools,
+        )
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model.py
index d28fa89c98241..f9fe332b0af7d 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model.py
@@ -26,11 +26,13 @@
 from datetime import datetime
 
 from vertexai.generative_models import HarmBlockThreshold, HarmCategory, Tool, grounding
+from vertexai.preview.evaluation import MetricPromptTemplateExamples
 
 from airflow.models.dag import DAG
 from airflow.providers.google.cloud.operators.vertex_ai.generative_model import (
     CountTokensOperator,
     GenerativeModelGenerateContentOperator,
+    RunEvaluationOperator,
     TextEmbeddingModelGetEmbeddingsOperator,
     TextGenerationModelPredictOperator,
 )
@@ -56,6 +58,37 @@
     HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
     HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
 }
+EVAL_DATASET = {
+    "context": [
+        "To make a classic spaghetti carbonara, start by bringing a large pot of salted water to a boil. While the water is heating up, cook pancetta or guanciale in a skillet with olive oil over medium heat until it's crispy and golden brown. Once the pancetta is done, remove it from the skillet and set it aside. In the same skillet, whisk together eggs, grated Parmesan cheese, and black pepper to make the sauce. When the pasta is cooked al dente, drain it and immediately toss it in the skillet with the egg mixture, adding a splash of the pasta cooking water to create a creamy sauce.",
+        "Preparing a perfect risotto requires patience and attention to detail. Begin by heating butter in a large, heavy-bottomed pot over medium heat. Add finely chopped onions and minced garlic to the pot, and cook until they're soft and translucent, about 5 minutes. Next, add Arborio rice to the pot and cook, stirring constantly, until the grains are coated with the butter and begin to toast slightly. Pour in a splash of white wine and cook until it's absorbed. From there, gradually add hot chicken or vegetable broth to the rice, stirring frequently, until the risotto is creamy and the rice is tender with a slight bite.",
+        "For a flavorful grilled steak, start by choosing a well-marbled cut of beef like ribeye or New York strip. Season the steak generously with kosher salt and freshly ground black pepper on both sides, pressing the seasoning into the meat. Preheat a grill to high heat and brush the grates with oil to prevent sticking. Place the seasoned steak on the grill and cook for about 4-5 minutes on each side for medium-rare, or adjust the cooking time to your desired level of doneness. Let the steak rest for a few minutes before slicing against the grain and serving.",
+        "Creating a creamy homemade tomato soup is a comforting and simple process. Begin by heating olive oil in a large pot over medium heat. Add diced onions and minced garlic to the pot and cook until they're soft and fragrant. Next, add chopped fresh tomatoes, chicken or vegetable broth, and a sprig of fresh basil to the pot. Simmer the soup for about 20-30 minutes, or until the tomatoes are tender and falling apart. Remove the basil sprig and use an immersion blender to puree the soup until smooth. Season with salt and pepper to taste before serving.",
+        "To bake a decadent chocolate cake from scratch, start by preheating your oven to 350°F (175°C) and greasing and flouring two 9-inch round cake pans. In a large mixing bowl, cream together softened butter and granulated sugar until light and fluffy. Beat in eggs one at a time, making sure each egg is fully incorporated before adding the next. In a separate bowl, sift together all-purpose flour, cocoa powder, baking powder, baking soda, and salt. Divide the batter evenly between the prepared cake pans and bake for 25-30 minutes, or until a toothpick inserted into the center comes out clean.",
+    ],
+    "instruction": ["Summarize the following article"] * 5,
+    "reference": [
+        "The process of making spaghetti carbonara involves boiling pasta, crisping pancetta or guanciale, whisking together eggs and Parmesan cheese, and tossing everything together to create a creamy sauce.",
+        "Preparing risotto entails sautéing onions and garlic, toasting Arborio rice, adding wine and broth gradually, and stirring until creamy and tender.",
+        "Grilling a flavorful steak involves seasoning generously, preheating the grill, cooking to desired doneness, and letting it rest before slicing.",
+        "Creating homemade tomato soup includes sautéing onions and garlic, simmering with tomatoes and broth, pureeing until smooth, and seasoning to taste.",
+        "Baking a decadent chocolate cake requires creaming butter and sugar, beating in eggs and alternating dry ingredients with buttermilk before baking until done.",
+    ],
+}
+METRICS = [
+    MetricPromptTemplateExamples.Pointwise.SUMMARIZATION_QUALITY,
+    MetricPromptTemplateExamples.Pointwise.GROUNDEDNESS,
+    MetricPromptTemplateExamples.Pointwise.VERBOSITY,
+    MetricPromptTemplateExamples.Pointwise.INSTRUCTION_FOLLOWING,
+    "exact_match",
+    "bleu",
+    "rouge_1",
+    "rouge_2",
+    "rouge_l_sum",
+]
+EXPERIMENT_NAME = "eval-experiment-airflow-operator"
+EXPERIMENT_RUN_NAME = "eval-experiment-airflow-operator-run"
+PROMPT_TEMPLATE = "{instruction}. Article: {context}. Summary:"
 
 with DAG(
     dag_id=DAG_ID,
@@ -108,6 +141,20 @@
     )
     # [END how_to_cloud_vertex_ai_generative_model_generate_content_operator]
 
+    # [START how_to_cloud_vertex_ai_run_evaluation_operator]
+    run_evaluation_task = RunEvaluationOperator(
+        task_id="run_evaluation_task",
+        project_id=PROJECT_ID,
+        location=REGION,
+        pretrained_model=MULTIMODAL_MODEL,
+        eval_dataset=EVAL_DATASET,
+        metrics=METRICS,
+        experiment_name=EXPERIMENT_NAME,
+        experiment_run_name=EXPERIMENT_RUN_NAME,
+        prompt_template=PROMPT_TEMPLATE,
+    )
+    # [END how_to_cloud_vertex_ai_run_evaluation_operator]
+
     from tests.system.utils.watcher import watcher
 
     # This test needs watcher in order to properly mark success/failure
