{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ur8xi4C7S06n"
   },
   "outputs": [],
   "source": [
    "# Copyright 2025 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JAPoU8Sm5E6e"
   },
   "source": [
    "# Get started with Vertex Prompt Optimizer - Custom metric\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fprompts%2Fprompt_optimizer%2Fget_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\">\n",
    "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\">\n",
    "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "</table>\n",
    "\n",
    "<div style=\"clear: both;\"></div>\n",
    "\n",
    "<b>Share to:</b>\n",
    "\n",
    "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
    "</a>\n",
    "\n",
    "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb\" target=\"_blank\">\n",
    "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
    "</a>            \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84f0f73a0f76"
   },
   "source": [
    "| Author(s) |\n",
    "| --- |\n",
    "| [Ivan Nardini](https://github.com/inardini) |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tvgnzT1CKxrO"
   },
   "source": [
    "## Overview\n",
    "\n",
    "When developing with large language models, crafting the perfect prompt—a process known as prompt engineering—is both an art and a science. It can be time-consuming and challenging to write prompts that consistently produce the desired results. Furthermore, as new and improved models are released, prompts that worked well before may need to be updated.\n",
    "\n",
    "To address these challenges, Vertex AI offers the **Prompt Optimizer**, a prompt optimization tool to help you refine and enhance your prompts automatically. This notebook serves as a comprehensive guide to both of its  approaches: the **Zero-Shot Optimizer** and the **Data-Driven Optimizer**.\n",
    "\n",
    "### The two approaches to prompt optimization\n",
    "\n",
    "#### 1\\. Zero-Shot Optimizer\n",
    "\n",
    "This is your go-to tool for rapid prompt refinement and generation *without* needing an evaluation dataset.\n",
    "\n",
    "  * **Generate from Scratch**: Simply describe a task in plain language, and it will generate a complete, well-structured system instruction for you.\n",
    "  * **Refine Existing Prompts**: Provide an existing prompt, and it will rewrite it based on established best practices for clarity, structure, and effectiveness.\n",
    "\n",
    "#### 2\\. Data-Driven Optimizer\n",
    "\n",
    "This tool performs a deep, performance-based optimization that uses your data to measure success.\n",
    "\n",
    "  * **Tune for Performance**: You provide a dataset of sample inputs and expected outputs, and it systematically tests and rewrites your system instructions to find the version that scores highest on the evaluation metrics you define.\n",
    "  * **Task-Specific**: It's the ideal choice when you want to fine-tune a prompt for a specific task and have data to prove what \"better\" looks like.\n",
    "\n",
    "In this tutorial, we'll show you how to optimize a simple prompt for a Gemini model using your own metric. The goal is to use Vertex AI Prompt Optimizer to find a new prompt template that generates better responses based on your own optimization metric."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U7rhvLMyJd2Y"
   },
   "source": [
    "## Get started\n",
    "\n",
    "Before we can start optimizing, we need to set up our Python environment and configure our Google Cloud project.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RGUH4MtEJd2Z"
   },
   "source": [
    "### Install required packages\n",
    "\n",
    "This command installs the necessary Python libraries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_FdPY3JMJd2Z"
   },
   "outputs": [],
   "source": [
    "%pip install \"google-cloud-aiplatform>=1.108.0\" \"pydantic\" \"etils\" \"protobuf==4.25.3\" --force-reinstall --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dmWOrTJ3gx13"
   },
   "source": [
    "### Authenticate your notebook environment (Colab only)\n",
    "\n",
    "If you are running this notebook in Google Colab, this cell handles authentication, allowing the notebook to securely access your Google Cloud resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NyKGtVQjgx13"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    from google.colab import auth\n",
    "\n",
    "    auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DF4l8DTdWgPY"
   },
   "source": [
    "### Set Google Cloud project information\n",
    "\n",
    "Here, we define essential variables for our Google Cloud project. The Prompt Optimizer job will run within a Google Cloud project. You need to [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com) and use the specified Cloud Storage bucket to read input data and write results.\n",
    "\n",
    "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Nqwi-5ufWp_B"
   },
   "outputs": [],
   "source": [
    "# Use the environment variable if the user doesn't provide Project ID.\n",
    "import os\n",
    "\n",
    "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
    "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
    "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
    "\n",
    "PROJECT_NUMBER = !gcloud projects describe {PROJECT_ID} --format=\"get(projectNumber)\"[0]\n",
    "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
    "\n",
    "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
    "\n",
    "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
    "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
    "\n",
    "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}\n",
    "\n",
    "import vertexai\n",
    "\n",
    "client = vertexai.Client(project=PROJECT_ID, location=LOCATION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AaksyUomxawt"
   },
   "source": [
    "### Service account and permissions\n",
    "\n",
    "The Prompt Optimizer runs as a backend job that needs permission to perform actions on your behalf. We grant the necessary IAM roles to the default Compute Engine service account, which the job uses to operate.\n",
    "\n",
    "  * `Vertex AI User`: Allows the job to call Vertex AI models.\n",
    "  * `Storage Object Admin`: Allows the job to read your dataset from and write results to your GCS bucket.\n",
    "  * `Artifact Registry Reader`: Allows the job to download necessary components.\n",
    "  * `Cloud Run Developer` to deploy function on Cloud Run.\n",
    "\n",
    "[Check out the documentation](https://cloud.google.com/iam/docs/manage-access-service-accounts#iam-view-access-sa-gcloud) to learn how to grant those permissions to a single service account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g7MNJEFP7-S9"
   },
   "outputs": [],
   "source": [
    "SERVICE_ACCOUNT = f\"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com\"\n",
    "\n",
    "for role in ['aiplatform.user', 'storage.objectAdmin', 'artifactregistry.reader', 'run.developer', 'run.invoker']:\n",
    "\n",
    "    ! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
    "      --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
    "      --role=roles/{role} --condition=None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5303c05f7aa6"
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6fc324893334"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import logging\n",
    "import subprocess\n",
    "from typing import Any, Dict, List, Optional, Tuple\n",
    "\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "from etils import epath\n",
    "import requests\n",
    "from google.cloud import storage\n",
    "from pydantic import BaseModel, Field\n",
    "logging.basicConfig(level=logging.INFO, force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3X-uCbdcN6MN"
   },
   "source": [
    "### Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qbz5ODl_OeAL"
   },
   "outputs": [],
   "source": [
    "def get_auth_token() -> str:\n",
    "    \"\"\"A function to collect the authorization token\"\"\"\n",
    "    result = subprocess.run(\n",
    "        [\"gcloud\", \"auth\", \"application-default\", \"print-access-token\", \"-q\"],\n",
    "        capture_output=True,\n",
    "        text=True,\n",
    "        check=True,\n",
    "    )\n",
    "    return result.stdout.strip()\n",
    "\n",
    "def format_demonstrations(demos: Any) -> List[str]:\n",
    "    \"\"\"Format demonstrations into readable strings.\"\"\"\n",
    "    if isinstance(demos, str):\n",
    "        try:\n",
    "            demos = json.loads(demos)\n",
    "        except (json.JSONDecodeError, ValueError):\n",
    "            return []\n",
    "\n",
    "    if not isinstance(demos, list):\n",
    "        return []\n",
    "\n",
    "    formatted = []\n",
    "    for demo in demos:\n",
    "        if isinstance(demo, dict):\n",
    "            demo_str = \"\\n\".join(f\"{k}: {v}\" for k, v in demo.items())\n",
    "            formatted.append(demo_str)\n",
    "        else:\n",
    "            formatted.append(str(demo))\n",
    "\n",
    "    return formatted\n",
    "\n",
    "\n",
    "def split_gcs_path(gcs_path: str) -> Tuple[str, str]:\n",
    "    \"\"\"Split GCS path into bucket name and prefix.\"\"\"\n",
    "    if not gcs_path.startswith(\"gs://\"):\n",
    "        raise ValueError(f\"Invalid GCS path. Must start with gs://: {gcs_path}\")\n",
    "\n",
    "    path = gcs_path[len(\"gs://\"):]\n",
    "    parts = path.split(\"/\", 1)\n",
    "    return parts[0], parts[1] if len(parts) > 1 else \"\"\n",
    "\n",
    "\n",
    "def list_gcs_objects(gcs_path: str) -> List[str]:\n",
    "    \"\"\"List all objects under given GCS path.\"\"\"\n",
    "    bucket_name, prefix = parse_gcs_path(gcs_path)\n",
    "\n",
    "    client = storage.Client()\n",
    "    bucket = client.bucket(bucket_name)\n",
    "    blobs = bucket.list_blobs(prefix=prefix)\n",
    "\n",
    "    return [blob.name for blob in blobs]\n",
    "\n",
    "\n",
    "def find_directories_with_files(\n",
    "    base_path: str, required_files: List[str]\n",
    ") -> List[str]:\n",
    "    \"\"\"Find directories containing all required files.\"\"\"\n",
    "    bucket_name, prefix = split_gcs_path(base_path)\n",
    "    all_paths = list_gcs_objects(base_path)\n",
    "\n",
    "    # Group files by directory\n",
    "    directories: Dict[str, set] = {}\n",
    "    for path in all_paths:\n",
    "        dir_path = \"/\".join(path.split(\"/\")[:-1])\n",
    "        filename = path.split(\"/\")[-1]\n",
    "\n",
    "        if dir_path not in directories:\n",
    "            directories[dir_path] = set()\n",
    "        directories[dir_path].add(filename)\n",
    "\n",
    "    # Find directories with all required files\n",
    "    matching_dirs = []\n",
    "    for dir_path, files in directories.items():\n",
    "        if all(req_file in files for req_file in required_files):\n",
    "            matching_dirs.append(f\"gs://{bucket_name}/{dir_path}\")\n",
    "\n",
    "    return matching_dirs\n",
    "\n",
    "def parse_gcs_path(gcs_path: str) -> Tuple[str, str]:\n",
    "    \"\"\"Parse GCS path into bucket name and prefix.\"\"\"\n",
    "    if not gcs_path.startswith(\"gs://\"):\n",
    "        raise ValueError(\"Invalid GCS path. Must start with gs://\")\n",
    "\n",
    "    path_without_prefix = gcs_path[5:]  # Remove 'gs://'\n",
    "    parts = path_without_prefix.split(\"/\", 1)\n",
    "    bucket_name = parts[0]\n",
    "    prefix = parts[1] if len(parts) > 1 else \"\"\n",
    "\n",
    "    return bucket_name, prefix\n",
    "\n",
    "def get_best_vapo_results(\n",
    "    base_path: str, metric_name: Optional[str] = None\n",
    ") -> Tuple[str, List[str]]:\n",
    "    \"\"\"Get the best system instruction and demonstrations across all VAPO runs.\"\"\"\n",
    "    # Find all valid runs\n",
    "    required_files = [\"eval_results.json\", \"templates.json\"]\n",
    "    runs = find_directories_with_files(base_path, required_files)\n",
    "\n",
    "    if not runs:\n",
    "        raise ValueError(f\"No valid runs found in {base_path}\")\n",
    "\n",
    "    best_score = float(\"-inf\")\n",
    "    best_instruction = \"\"\n",
    "    best_demonstrations: List[str] = []\n",
    "\n",
    "    for run_path in runs:\n",
    "        try:\n",
    "            # Check main templates.json first\n",
    "            templates_path = f\"{run_path}/templates.json\"\n",
    "            with epath.Path(templates_path).open(\"r\") as f:\n",
    "                templates_data = json.load(f)\n",
    "\n",
    "            if templates_data:\n",
    "                df = pd.json_normalize(templates_data)\n",
    "\n",
    "                # Find metric column\n",
    "                metric_columns = [\n",
    "                    col for col in df.columns\n",
    "                    if \"metric\" in col and \"mean\" in col\n",
    "                ]\n",
    "\n",
    "                if metric_columns:\n",
    "                    # Select appropriate metric\n",
    "                    if metric_name:\n",
    "                        metric_col = next(\n",
    "                            (col for col in metric_columns if metric_name in col),\n",
    "                            None\n",
    "                        )\n",
    "                    else:\n",
    "                        composite_cols = [\n",
    "                            col for col in metric_columns\n",
    "                            if \"composite_metric\" in col\n",
    "                        ]\n",
    "                        metric_col = (\n",
    "                            composite_cols[0] if composite_cols else metric_columns[0]\n",
    "                        )\n",
    "\n",
    "                    if metric_col and metric_col in df.columns:\n",
    "                        best_idx = df[metric_col].argmax()\n",
    "                        score = float(df.iloc[best_idx][metric_col])\n",
    "\n",
    "                        if score > best_score:\n",
    "                            best_score = score\n",
    "                            best_row = df.iloc[best_idx]\n",
    "\n",
    "                            # Extract instruction if present\n",
    "                            if \"prompt\" in best_row or \"instruction\" in best_row:\n",
    "                                instruction = best_row.get(\n",
    "                                    \"prompt\", best_row.get(\"instruction\", \"\")\n",
    "                                )\n",
    "                                if instruction:\n",
    "                                    instruction = instruction.replace(\n",
    "                                        \"store('answer', llm())\", \"{{llm()}}\"\n",
    "                                    )\n",
    "                                    best_instruction = instruction\n",
    "\n",
    "                            # Extract demonstrations if present\n",
    "                            if \"demonstrations\" in best_row or \"demo_set\" in best_row:\n",
    "                                demos = best_row.get(\n",
    "                                    \"demonstrations\", best_row.get(\"demo_set\", [])\n",
    "                                )\n",
    "                                best_demonstrations = format_demonstrations(demos)\n",
    "\n",
    "            # Check instruction-specific optimization\n",
    "            instruction_path = f\"{run_path}/instruction/templates.json\"\n",
    "            try:\n",
    "                with epath.Path(instruction_path).open(\"r\") as f:\n",
    "                    instruction_data = json.load(f)\n",
    "\n",
    "                if instruction_data:\n",
    "                    inst_df = pd.json_normalize(instruction_data)\n",
    "                    metric_columns = [\n",
    "                        col for col in inst_df.columns\n",
    "                        if \"metric\" in col and \"mean\" in col\n",
    "                    ]\n",
    "\n",
    "                    if metric_columns:\n",
    "                        if metric_name:\n",
    "                            metric_col = next(\n",
    "                                (col for col in metric_columns if metric_name in col),\n",
    "                                None,\n",
    "                            )\n",
    "                        else:\n",
    "                            composite_cols = [\n",
    "                                col for col in metric_columns\n",
    "                                if \"composite_metric\" in col\n",
    "                            ]\n",
    "                            metric_col = (\n",
    "                                composite_cols[0] if composite_cols else metric_columns[0]\n",
    "                            )\n",
    "\n",
    "                        if metric_col and metric_col in inst_df.columns:\n",
    "                            inst_best_idx = inst_df[metric_col].argmax()\n",
    "                            inst_score = float(inst_df.iloc[inst_best_idx][metric_col])\n",
    "\n",
    "                            if inst_score > best_score:\n",
    "                                best_score = inst_score\n",
    "                                best_row = inst_df.iloc[inst_best_idx]\n",
    "\n",
    "                                instruction = best_row.get(\n",
    "                                    \"prompt\", best_row.get(\"instruction\", \"\")\n",
    "                                )\n",
    "                                if instruction:\n",
    "                                    instruction = instruction.replace(\n",
    "                                        \"store('answer', llm())\", \"{{llm()}}\"\n",
    "                                    )\n",
    "                                    best_instruction = instruction\n",
    "                                # In instruction-only mode, there might not be demonstrations\n",
    "                                if \"demonstrations\" not in best_row and \"demo_set\" not in best_row:\n",
    "                                    best_demonstrations = []\n",
    "            except FileNotFoundError:\n",
    "                pass\n",
    "\n",
    "            # Check demonstration-specific optimization\n",
    "            demo_path = f\"{run_path}/demonstration/templates.json\"\n",
    "            try:\n",
    "                with epath.Path(demo_path).open(\"r\") as f:\n",
    "                    demo_data = json.load(f)\n",
    "\n",
    "                if demo_data:\n",
    "                    demo_df = pd.json_normalize(demo_data)\n",
    "                    metric_columns = [\n",
    "                        col for col in demo_df.columns\n",
    "                        if \"metric\" in col and \"mean\" in col\n",
    "                    ]\n",
    "\n",
    "                    if metric_columns:\n",
    "                        if metric_name:\n",
    "                            metric_col = next(\n",
    "                                (col for col in metric_columns if metric_name in col),\n",
    "                                None,\n",
    "                            )\n",
    "                        else:\n",
    "                            composite_cols = [\n",
    "                                col for col in metric_columns\n",
    "                                if \"composite_metric\" in col\n",
    "                            ]\n",
    "                            metric_col = (\n",
    "                                composite_cols[0] if composite_cols else metric_columns[0]\n",
    "                            )\n",
    "\n",
    "                        if metric_col and metric_col in demo_df.columns:\n",
    "                            demo_best_idx = demo_df[metric_col].argmax()\n",
    "                            demo_score = float(demo_df.iloc[demo_best_idx][metric_col])\n",
    "\n",
    "                            if demo_score > best_score:\n",
    "                                best_score = demo_score\n",
    "                                best_row = demo_df.iloc[demo_best_idx]\n",
    "\n",
    "                                demos = best_row.get(\n",
    "                                    \"demonstrations\", best_row.get(\"demo_set\", [])\n",
    "                                )\n",
    "                                best_demonstrations = format_demonstrations(demos)\n",
    "                                # In demo-only mode, there might not be an instruction\n",
    "                                if \"prompt\" not in best_row and \"instruction\" not in best_row:\n",
    "                                    best_instruction = \"\"\n",
    "                                else:\n",
    "                                    instruction = best_row.get(\n",
    "                                        \"prompt\", best_row.get(\"instruction\", \"\")\n",
    "                                    )\n",
    "                                    if instruction:\n",
    "                                        instruction = instruction.replace(\n",
    "                                            \"store('answer', llm())\", \"{{llm()}}\"\n",
    "                                        )\n",
    "                                        best_instruction = instruction\n",
    "            except (FileNotFoundError, json.JSONDecodeError):\n",
    "                pass\n",
    "\n",
    "        except Exception as e:\n",
    "            logging.warning(f\"Error processing run {run_path}: {e}\")\n",
    "            continue\n",
    "\n",
    "    if best_score == float(\"-inf\"):\n",
    "        raise ValueError(\"Could not find any valid results\")\n",
    "\n",
    "    return best_instruction, best_demonstrations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FUr06qOWxuy9"
   },
   "source": [
    "## Using the Data-Driven Optimizer approach with your own custom metrics\n",
    "\n",
    "The following sections will guide you through setting up your environment, preparing your data, deploy the custom metrics, and running an optimization job to find a better prompt using the data-driven optimizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tfpGmIWrVEt1"
   },
   "source": [
    "### Preparing the Data and Running the Job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The prompt template to optimize\n",
    "\n",
    "A prompt consists of two key parts:\n",
    "\n",
    "* **System Instruction Template** which is a fixed part of the prompt that control or alter the model's behavior across all queries for a given task.\n",
    "\n",
    "* **Prompt Template** which is a dynamic part of the prompt that changes based on the task. Prompt template includes context, task and more. To learn more, see [components of a prompt](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/prompts/prompt-design-strategies#components-of-a-prompt) in the official documentation.\n",
    "\n",
    "In this scenario, you use Vertex AI prompt optimizer to optimize a simple system instruction template. And you use some examples in the remaining prompt template for evaluating different instruction templates along the optimization process.\n",
    "\n",
    "> Important: Having the `target` placeholder in the prompt template is optional. It represents the prompt's ground truth response in your prompt optimization dataset that you aim to optimize for your templates. If you don't have the prompt's ground truth response, remember to set the `source_model` parameter to your prompt optimizer configuration (see below) instead of adding ground truth responses. Vertex AI prompt optimizer would run your sample prompts on the source model to generate the ground truth responses for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_instruction = \"\"\"\n",
    "Given a question with context, provide the correct answer to the question.\n",
    "\"\"\"\n",
    "\n",
    "prompt_template = \"\"\"\n",
    "Some examples of correct answer to a question are:\n",
    "Question: {question}\n",
    "Context: {ctx}\n",
    "Answer: {target}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0xXys6YmDS2v"
   },
   "source": [
    "#### The optimization dataset\n",
    "\n",
    "The optimizer's performance depends heavily on the quality of your sample data.\n",
    "\n",
    "For this example, we use a question-answering dataset where each row contains a `question`, context (`ctx`), and a ground-truth `target` answer. The `{target}` variable is crucial for computation-based evaluation metrics like `question_answering_correctness`.\n",
    "\n",
    "> Important: For effective **prompt optimization**, provide a dataset of examples where your model is poor in performance when using current system instruction template. For reliable results, use 50-100 distinct samples. In case of **prompt migration**, consider using the source model to label examples that the target model struggles with, helping to identify areas for improvement.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QWMSgAdWDWwW"
   },
   "outputs": [],
   "source": [
    "input_data_path = \"gs://github-repo/prompts/prompt_optimizer/rag_qa_dataset.jsonl\"\n",
    "prompt_optimization_df = pd.read_json(input_data_path, lines=True)\n",
    "prompt_optimization_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hxpid3KgAkYM"
   },
   "source": [
    "#### Define and deploy your own custom optimization metric on Cloud function\n",
    "\n",
    "To optimize your prompt template using a custom optimization metric, you need to deploy a function with your own metric code on a Cloud function. To deploy a Cloud function with your own custom metric, you cover the following steps:\n",
    "\n",
    "1.   Define requirements\n",
    "2.   Write your own custom metric function code\n",
    "3.   Deploy the custom code as Cloud function\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nxh2e88fAnQc"
   },
   "source": [
    "##### Define requirements\n",
    "\n",
    "Set the custom metric dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q-hUlhgBCus4"
   },
   "outputs": [],
   "source": [
    "build_path = Path(\"build\")\n",
    "build_path.mkdir(exist_ok=True)\n",
    "\n",
    "requirements = \"\"\"\n",
    "functions-framework==3.*\n",
    "google-cloud-aiplatform>=1.108.0\n",
    "\"\"\"\n",
    "\n",
    "with open(build_path / \"requirements.txt\", \"w\") as f:\n",
    "    f.write(requirements)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k_EFZEBeEy48"
   },
   "source": [
    "##### Write your own custom metric function\n",
    "\n",
    "Define the module which contains your own custom metric function definition.\n",
    "\n",
    "In this case, you have a custom evaluation metric to evaluate the user engagement and personalization. The custom evaluation metric is defined using the `evaluate_engagement_personalization_fn`.\n",
    "\n",
    "The function leverages Gemini 2.5 Flash to act as an LLM-as-Judge. It sends a prompt to the auto-rater, receives a score (1-5), and an explanation, then returns these as a dictionary containing two fields: the custom metric's score (as you defined it) and an explanation of how this metric helps optimize the prompt template.\n",
    "\n",
    "You use the `main` function to deploy the `evaluate_engagement_personalization_fn` function as a Cloud Function, receiving a question, response, and a target response as input and returning the auto-rater's evaluation.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1wGoVQNCMbxe"
   },
   "outputs": [],
   "source": [
    "custom_metric_function_code = '''\n",
    "\"\"\"\n",
    "This module contains the custom evaluation metric definition to optimize a prompt template with Vertex AI prompt optimizer\n",
    "\"\"\"\n",
    "\n",
    "from typing import Dict\n",
    "from vertexai.generative_models import (\n",
    "    GenerationConfig,\n",
    "    GenerativeModel,\n",
    "    HarmBlockThreshold,\n",
    "    HarmCategory,\n",
    ")\n",
    "\n",
    "import json\n",
    "import functions_framework\n",
    "\n",
    "def get_autorater_response(metric_prompt: str) -> dict:\n",
    "    \"\"\"This function is to generate the evaluation response from the autorater.\"\"\"\n",
    "\n",
    "    metric_response_schema = {\n",
    "        \"type\": \"OBJECT\",\n",
    "        \"properties\": {\n",
    "            \"score\": {\"type\": \"NUMBER\"},\n",
    "            \"explanation\": {\"type\": \"STRING\"},\n",
    "        },\n",
    "        \"required\": [\"score\", \"explanation\"],\n",
    "    }\n",
    "\n",
    "    autorater = GenerativeModel(\n",
    "        \"gemini-2.5-flash\",\n",
    "        generation_config=GenerationConfig(\n",
    "            response_mime_type=\"application/json\",\n",
    "            response_schema=metric_response_schema,\n",
    "        ),\n",
    "        safety_settings={\n",
    "            HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,\n",
    "            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,\n",
    "            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,\n",
    "            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,\n",
    "            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,\n",
    "        },\n",
    "    )\n",
    "\n",
    "    response = autorater.generate_content(metric_prompt)\n",
    "\n",
    "    response_json = {}\n",
    "    response_json = json.loads(response.text)\n",
    "    return response_json\n",
    "\n",
    "\n",
    "# Define custom evaluation criteria\n",
    "def evaluate_engagement_personalization_fn(question: str, response:str, target: str) -> Dict[str, str]:\n",
    "    \"\"\"Evaluates an AI-generated response for User Engagement and Personalization.\"\"\"\n",
    "\n",
    "    custom_metric_prompt_template = \"\"\"\n",
    "\n",
    "    # Instruction\n",
    "    You are an expert evaluator. Your task is to evaluate the quality of the LLM-generated responses against a reference target response.\n",
    "    You should first read the Question carefully, and then evaluate the quality of the responses based on the Criteria provided in the Evaluation section below.\n",
    "    You will assign the response a rating following the Rating Rubric only and a step-by-step explanation for your rating.\n",
    "\n",
    "    # Evaluation\n",
    "\n",
    "    ## Criteria\n",
    "    Relevance and Customization: The response should directly address the user's query and demonstrate an understanding of their specific needs or preferences, such as dietary restrictions, skill level, or taste preferences.\n",
    "    Interactivity and Proactiveness: The response should go beyond simply answering the question by actively encouraging further interaction through follow-up questions, suggestions for additional exploration, or prompts for more information to provide a tailored experience.\n",
    "    Tone and Empathy: The response should adopt an appropriate and empathetic tone that fosters a positive and supportive user experience, making the user feel heard and understood.\n",
    "\n",
    "    ## Rating rubric\n",
    "    1 - Minimal: The response lacks personalization and demonstrates minimal engagement with the user. The tone may be impersonal or generic.\n",
    "    2 - Basic: The response shows some basic personalization but lacks depth or specificity. Engagement is limited, possibly with generic prompts or suggestions. The tone is generally neutral but may lack warmth or empathy.\n",
    "    3 - Moderate: The response demonstrates clear personalization and attempts to engage the user with relevant follow-up questions or prompts based on their query. The tone is friendly and supportive, fostering a positive user experience.\n",
    "    4 - High: The response demonstrates a high degree of personalization and actively engages the user with relevant follow-up questions or prompts. The tone is empathetic and understanding, creating a strong connection with the user.\n",
    "    5 - Exceptional: The response goes above and beyond to personalize the experience, anticipating user needs, and fostering a genuine connection. The tone is warm, encouraging, and inspiring, leaving the user feeling empowered and motivated.\n",
    "\n",
    "    ## Evaluation steps\n",
    "    Step 1: Carefully read both the question and the generated response. Ensure a clear understanding of the user's intent, needs, and any specific context provided.\n",
    "    Step 2: Evaluate how well the response directly addresses the user's query and demonstrates an understanding of their specific needs or preferences.\n",
    "    Step 3: Determine the extent to which the response actively encourages further interaction and provides a tailored experience.\n",
    "    Step 4: Evaluate Tone & Empathy: Analyze the tone of the response, ensuring it fosters a positive and supportive user experience, making the user feel heard and understood.\n",
    "    Step 5: Based on the three criteria above, assign a score from 1 to 5 according to the score rubric.\n",
    "    Step 6: Justify the assigned score with a clear and concise explanation, highlighting the strengths and weaknesses of the response with respect to each criterion.\n",
    "\n",
    "    # Question : {question}\n",
    "    # Generated response: {response}\n",
    "    # Reference response: {target}\n",
    "    \"\"\"\n",
    "\n",
    "    custom_metric_prompt = custom_metric_prompt_template.format(question=question, response=response, target=target)\n",
    "    response_dict = get_autorater_response(custom_metric_prompt)\n",
    "\n",
    "    return {\n",
    "        \"custom_engagement_personalization_score\": response_dict[\"score\"],\n",
    "        \"explanation\": response_dict[\"explanation\"],\n",
    "    }\n",
    "\n",
    "# Register an HTTP function with the Functions Framework\n",
    "@functions_framework.http\n",
    "def main(request):\n",
    "  request_json = request.get_json(silent=True)\n",
    "\n",
    "  if not request_json:\n",
    "    raise ValueError('Cannot find request json.')\n",
    "\n",
    "  question = request_json['question']\n",
    "  response = request_json['response']\n",
    "  reference = request_json['target']\n",
    "\n",
    "  get_evaluation_result = evaluate_engagement_personalization_fn(question, response, reference)\n",
    "  return json.dumps(get_evaluation_result)\n",
    "'''\n",
    "\n",
    "with open(build_path / \"main.py\", \"w\") as f:\n",
    "    f.write(custom_metric_function_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T7R0LDZMCPnL"
   },
   "source": [
    "##### Deploy the custom metric as a Cloud Function\n",
    "\n",
    "Use gcloud command line to deploy a Cloud function. To learn more, check out [Deploy a Python service to Cloud Run](https://cloud.google.com/run/docs/quickstarts/build-and-deploy/deploy-python-service) quickstart.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nwBZGvkLCizs"
   },
   "outputs": [],
   "source": [
    "!gcloud functions deploy 'custom_engagement_personalization_metric' \\\n",
    " --gen2 \\\n",
    " --runtime=\"python310\" \\\n",
    " --source={str(build_path)} \\\n",
    " --entry-point=main \\\n",
    " --trigger-http \\\n",
    " --timeout=3600 \\\n",
    " --memory=2Gb \\\n",
    " --concurrency=6 \\\n",
    " --min-instances=6 \\\n",
    " --project {PROJECT_ID} \\\n",
    " --region={LOCATION} \\\n",
    " --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FoEypczSwAGK"
   },
   "source": [
    "##### Test your custom evaluation metric\n",
    "\n",
    "After you deploy your  custom evaluation metric as Cloud function, submit a request to validate the output of the custom evaluation function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HXOWYp2MwEsA"
   },
   "outputs": [],
   "source": [
    "custom_evaluator_function_uri = ! gcloud functions describe 'custom_engagement_personalization_metric' --gen2 --region {LOCATION} --format=\"value(url)\"\n",
    "custom_evaluator_function_uri = custom_evaluator_function_uri[0].strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0JMeIyx0DHnc"
   },
   "outputs": [],
   "source": [
    "headers = {\n",
    "    \"Authorization\": f\"Bearer {get_auth_token()}\",\n",
    "    \"Content-Type\": \"application/json\",\n",
    "}\n",
    "\n",
    "json_data = {\n",
    "    \"question\": \"\"\"\n",
    "      What are some techniques for cooking red meat and pork that maximize flavor and tenderness while minimizing the formation of unhealthy compounds?\n",
    "      \"\"\",\n",
    "    \"response\": \"\"\"\n",
    "      * Marinating in acidic ingredients like lemon juice or vinegar to tenderize the meat \\n * Cooking to an internal temperature of 145°F (63°C) for safety \\n * Using high-heat cooking methods like grilling and pan-searing for browning and caramelization \\n * Avoiding charring to minimize the formation of unhealthy compounds\n",
    "      \"\"\",\n",
    "    \"target\": \"\"\"\n",
    "      Here's how to tackle those delicious red meats and pork while keeping things healthy:\n",
    "      **Prioritize Low and Slow:**\n",
    "      * **Braising and Stewing:** These techniques involve gently simmering meat in liquid over low heat for an extended period. This breaks down tough collagen, resulting in incredibly tender and flavorful meat. Plus, since the cooking temperature is lower, it minimizes the formation of potentially harmful compounds associated with high-heat cooking.\n",
    "      * **Sous Vide:** This method involves sealing meat in a vacuum bag and immersing it in a precisely temperature-controlled water bath.  It allows for even cooking to the exact desired doneness, resulting in incredibly juicy and tender meat.  Because the temperature is controlled and lower than traditional methods, it can be a healthier option.\n",
    "      **High Heat Tips:**\n",
    "      * **Marinades are Your Friend:** As you mentioned, acidic marinades tenderize meat.  They also add flavor!\n",
    "      * **Temperature Control is Key:**  Use a meat thermometer to ensure you reach the safe internal temperature of 145°F (63°C) without overcooking.\n",
    "      * **Don't Burn It!**  While some browning is desirable, charring creates those unhealthy compounds.  Pat meat dry before cooking to minimize steaming and promote browning.  Let the pan heat up properly before adding the meat to achieve a good sear.\n",
    "\n",
    "      **Remember:**  Trim visible fat before cooking to reduce saturated fat content.  Let meat rest after cooking; this allows juices to redistribute, resulting in a more tender and flavorful final product.\n",
    "      \"\"\",\n",
    "}\n",
    "\n",
    "response = requests.post(\n",
    "    custom_evaluator_function_uri, headers=headers, json=json_data, timeout=70\n",
    ").json()\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ceZNbD_YzLEY"
   },
   "source": [
    "#### Set optimization configuration\n",
    "\n",
    "Now, we'll create a dictionary with our specific settings and use it to instantiate our `OptimizationConfig` class.\n",
    "\n",
    "The `OptimizationConfig` class, built using `pydantic`, acts as a structured and validated blueprint for our optimization task. It ensures all necessary parameters are defined before we submit the job.\n",
    "\n",
    "In this scenario, you have two additional parameters:\n",
    "\n",
    "* `custom_metric_name` parameter which lets you pass your own custom metric to optimizer the prompt template.\n",
    "\n",
    "* `custom_metric_cloud_function_name` parameter which indicates the Cloud function to call for collecting custom function evaluation metric output.\n",
    "\n",
    "For more advanced control, you can learn and explore more about all the parameters and how to best use them in the [detailed documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/prompts/data-driven-optimizer).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5-yioYmFCBEB"
   },
   "outputs": [],
   "source": [
    "class OptimizationConfig(BaseModel):\n",
    "    \"\"\"\n",
    "    A comprehensive prompt optimization configuration model.\n",
    "    \"\"\"\n",
    "\n",
    "    # Basic Configuration\n",
    "    system_instruction: str = Field(\n",
    "        ...,\n",
    "        description=\"System instructions for the target model. String. This field is required.\",\n",
    "    )\n",
    "    prompt_template: str = Field(\n",
    "        ..., description=\"Template for prompts. String. This field is required.\"\n",
    "    )\n",
    "    target_model: str = Field(\n",
    "        \"gemini-2.5-flash\",\n",
    "        description='Target model for optimization. Supported models: \"gemini-2.5-flash\", \"gemini-2.5-pro\"',\n",
    "    )\n",
    "    thinking_budget: int = Field(\n",
    "        -1,\n",
    "        description=\"Thinking budget for thinking models. -1 means auto/no thinking. Integer.\",\n",
    "    )\n",
    "    optimization_mode: str = Field(\n",
    "        \"instruction\",\n",
    "        description='Optimization mode. Supported modes: \"instruction\", \"demonstration\", \"instruction_and_demo\".',\n",
    "    )\n",
    "    project: str = Field(\n",
    "        ..., description=\"Google Cloud project ID. This field is required.\"\n",
    "    )\n",
    "\n",
    "    # Evaluation Settings\n",
    "    eval_metrics_types: List[str] = Field(\n",
    "        description='List of evaluation metrics. E.g., \"bleu\", \"rouge_l\", \"safety\".'\n",
    "    )\n",
    "    eval_metrics_weights: List[float] = Field(\n",
    "        description=\"Weights for evaluation metrics. Length must match eval_metrics_types and should sum to 1.\"\n",
    "    )\n",
    "    aggregation_type: str = Field(\n",
    "        \"weighted_sum\",\n",
    "        description='Aggregation type for metrics. Supported: \"weighted_sum\", \"weighted_average\".',\n",
    "    )\n",
    "    custom_metric_name: str = Field(\n",
    "        \"\",\n",
    "        description=\"Metric name, as defined by the key that corresponds in the dictionary returned from Cloud function. String.\",\n",
    "    )\n",
    "    custom_metric_cloud_function_name: str = Field(\n",
    "        \"\",\n",
    "        description=\"Cloud Run function name you previously deployed. String.\",\n",
    "    )\n",
    "\n",
    "    # Data and I/O Paths\n",
    "    input_data_path: str = Field(\n",
    "        ...,\n",
    "        description=\"Cloud Storage URI to input optimization data. This field is required.\",\n",
    "    )\n",
    "    output_path: str = Field(\n",
    "        ...,\n",
    "        description=\"Cloud Storage URI to save optimization results. This field is required.\",\n",
    "    )\n",
    "\n",
    "    # (Optional) Advanced Configuration\n",
    "    num_steps: int = Field(\n",
    "        10,\n",
    "        ge=10,\n",
    "        le=20,\n",
    "        description=\"Number of iterations in instruction optimization mode. Integer between 10 and 20.\",\n",
    "    )\n",
    "    num_demo_set_candidates: int = Field(\n",
    "        10,\n",
    "        ge=10,\n",
    "        le=30,\n",
    "        description=\"Number of demonstrations evaluated. Integer between 10 and 30.\",\n",
    "    )\n",
    "    demo_set_size: int = Field(\n",
    "        3,\n",
    "        ge=3,\n",
    "        le=6,\n",
    "        description=\"Number of demonstrations generated per prompt. Integer between 3 and 6.\",\n",
    "    )\n",
    "\n",
    "    # (Optional) Model Locations and QPS\n",
    "    target_model_location: str = Field(\n",
    "        \"us-central1\", description=\"Location of the target model. Default us-central1.\"\n",
    "    )\n",
    "    target_model_qps: int = Field(\n",
    "        1,\n",
    "        ge=1,\n",
    "        description=\"QPS for the target model. Integer >= 1, based on your quota.\",\n",
    "    )\n",
    "    optimizer_model_location: str = Field(\n",
    "        \"us-central1\",\n",
    "        description=\"Location of the optimizer model. Default us-central1.\",\n",
    "    )\n",
    "    optimizer_model_qps: int = Field(\n",
    "        1,\n",
    "        ge=1,\n",
    "        description=\"QPS for the optimization model. Integer >= 1, based on your quota.\",\n",
    "    )\n",
    "    source_model: str = Field(\n",
    "        \"\",\n",
    "        description=\"Google model previously used with these prompts. Not needed if providing a target column.\",\n",
    "    )\n",
    "    source_model_location: str = Field(\n",
    "        \"us-central1\", description=\"Location of the source model. Default us-central1.\"\n",
    "    )\n",
    "    source_model_qps: Optional[int] = Field(\n",
    "        None, ge=1, description=\"Optional QPS for the source model. Integer >= 1.\"\n",
    "    )\n",
    "    eval_qps: int = Field(\n",
    "        1,\n",
    "        ge=1,\n",
    "        description=\"QPS for the eval model. Integer >= 1, based on your quota.\",\n",
    "    )\n",
    "\n",
    "    # (Optional) Response, Language, and Data Handling\n",
    "    response_mime_type: str = Field(\n",
    "        \"text/plain\",\n",
    "        description=\"MIME response type from the target model. E.g., 'text/plain', 'application/json'.\",\n",
    "    )\n",
    "    response_schema: str = Field(\n",
    "        \"\", description=\"The Vertex AI Controlled Generation response schema.\"\n",
    "    )\n",
    "    language: str = Field(\n",
    "        \"English\",\n",
    "        description='Language of the system instructions. E.g., \"English\", \"Japanese\".',\n",
    "    )\n",
    "    placeholder_to_content: Dict[str, Any] = Field(\n",
    "        {},\n",
    "        description=\"Dictionary of placeholders to replace parameters in the system instruction.\",\n",
    "    )\n",
    "    data_limit: int = Field(\n",
    "        10,\n",
    "        ge=5,\n",
    "        le=100,\n",
    "        description=\"Amount of data used for validation. Integer between 5 and 100.\",\n",
    "    )\n",
    "    translation_source_field_name: str = Field(\n",
    "        \"\",\n",
    "        description=\"Field name for source text if using translation metrics (Comet, MetricX).\",\n",
    "    )\n",
    "    has_multimodal_inputs: bool = Field(\n",
    "        False, description=\"Whether the input data is multimodal.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZBimPdKrbUZy"
   },
   "source": [
    "Set the optimization configuration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "40Pyzkot040M"
   },
   "outputs": [],
   "source": [
    "output_path = f\"{BUCKET_URI}/optimization_results/\"\n",
    "\n",
    "vapo_data_settings = {\n",
    "    \"system_instruction\": system_instruction,\n",
    "    \"prompt_template\": prompt_template,\n",
    "    \"target_model\": \"gemini-2.5-flash\",\n",
    "    \"thinking_budget\": -1,\n",
    "    \"optimization_mode\": \"instruction\",\n",
    "    \"custom_metric_name\": \"custom_engagement_personalization_score\",\n",
    "    \"custom_metric_cloud_function_name\": \"custom_engagement_personalization_metric\",\n",
    "    \"eval_metrics_types\": [\"question_answering_correctness\", \"custom_metric\"],\n",
    "    \"eval_metrics_weights\": [0.8, 0.2],\n",
    "    \"aggregation_type\": \"weighted_sum\",\n",
    "    \"input_data_path\": input_data_path,\n",
    "    \"output_path\": output_path,\n",
    "    \"project\": PROJECT_ID,\n",
    "}\n",
    "\n",
    "vapo_data_config = OptimizationConfig(**vapo_data_settings)\n",
    "vapo_data_config_json = vapo_data_config.model_dump()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "92qSHhIT838O"
   },
   "source": [
    "#### Upload configuration to Cloud Storage\n",
    "\n",
    "Write the Prompt Optimizer configuration to the file in your GCS bucket.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6PG_a6ss4J1l"
   },
   "outputs": [],
   "source": [
    "config_path = f\"{BUCKET_URI}/config.json\"\n",
    "\n",
    "with epath.Path(config_path).open(\"w\") as config_file:\n",
    "    json.dump(vapo_data_config_json, config_file)\n",
    "config_file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FpRGZTk68-Nu"
   },
   "source": [
    "#### Run the prompt optimization job\n",
    "\n",
    "This is the final step. We pass the path to our configuration file and the service account to the Vertex AI client. The `optimize` method starts the custom job on the Vertex AI backend. We set `wait_for_completion` to `True` so the script will pause until the job is finished.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uGZKNjsu6EEw"
   },
   "outputs": [],
   "source": [
    "vapo_data_run_config = {\n",
    "    \"config_path\": config_path,\n",
    "    \"wait_for_completion\": True,\n",
    "    \"service_account\": SERVICE_ACCOUNT,\n",
    "}\n",
    "\n",
    "result = client.prompt_optimizer.optimize(method=\"vapo\", config=vapo_data_run_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wJL6tRAWKyXz"
   },
   "source": [
    "### Get and use the best prompt programmatically\n",
    "\n",
    "For use in an application, you can programmatically retrieve the top-performing instruction from the output files stored in GCS.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8b_LRAhyxOvQ"
   },
   "outputs": [],
   "source": [
    "best_instruction, _ = get_best_vapo_results(output_path)\n",
    "print(\"The optimized instruction is:\\n\", best_instruction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2a4e033321ad"
   },
   "source": [
    "## Cleaning up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U_n-_B-Ekpk_"
   },
   "outputs": [],
   "source": [
    "delete_job = True\n",
    "delete_bucket = True\n",
    "delete_function = True\n",
    "\n",
    "if delete_job:\n",
    "    from google.cloud import aiplatform\n",
    "    aiplatform.init(project=PROJECT_ID, location=LOCATION)\n",
    "    custom_job_list = aiplatform.CustomJob.list()\n",
    "    if custom_job_list:\n",
    "        latest_job = custom_job_list[0]\n",
    "        latest_job.delete()\n",
    "\n",
    "if delete_function:\n",
    "    !gcloud functions delete 'custom_engagement_personalization_metric' --gen2 --region {LOCATION} --quiet\n",
    "\n",
    "if delete_bucket:\n",
    "    ! gsutil -m rm -r $BUCKET_URI"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb",
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": "workbench-notebooks.m131",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m131"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
