{
 "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 AI Prompt Optimizer\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.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.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.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.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.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.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.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.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.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>"
   ]
  },
  {
   "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 will walk through both methods. First, we'll explore the **Zero-Shot Optimizer** for quick, data-free improvements. Then, we'll dive deep into the **Data-Driven Optimizer**, learning how to leverage a dataset to achieve the best possible performance for a specific task.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "61RBz8LLbxCR"
   },
   "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": "No17Cw5hgx12"
   },
   "source": [
    "### Install required packages\n",
    "\n",
    "This command installs the necessary Python libraries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tFy3H3aPgx12"
   },
   "outputs": [],
   "source": [
    "%pip install \"google-cloud-aiplatform>=1.108.0\" \"pydantic\" \"etils\" \"protobuf==4.25.3\" \"gradio\" --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",
    "\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']:\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 io\n",
    "import json\n",
    "import logging\n",
    "import re\n",
    "from typing import Any, Dict, List, Optional, Tuple\n",
    "\n",
    "import gradio as gr\n",
    "import pandas as pd\n",
    "from etils import epath\n",
    "from google.cloud import storage\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "logging.basicConfig(level=logging.INFO, force=True)\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9EfCy5RI19vt"
   },
   "source": [
    "### Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PCIt8uAxYZKZ"
   },
   "outputs": [],
   "source": [
    "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": "mL5FRIvTTDfk"
   },
   "source": [
    "## **Part 1: Zero-Shot Optimizer**\n",
    "\n",
    "We'll begin with the zero-shot approach. The following section will guide you through the process of optimizing your prompt without providing additional examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y3Bp8o8NAO04"
   },
   "source": [
    "### Run a Zero-shot optimization job\n",
    "\n",
    "To run a `Zero-shot optimization job`, you can use the `optimize_prompt` method. The service will use a research-based metaprompt to optimize your initial prompt.\n",
    "\n",
    "**Notice:** The response also provides a more detailed output with `.applicable_guidelines`. This information is useful for understanding how your prompt was improved and possibly help you writing better prompts.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VLLA617pbuHg"
   },
   "outputs": [],
   "source": [
    "prompt = \"You are a helpful assistant. Given a question with context, provide the correct answer to the question.\"\n",
    "response = client.prompt_optimizer.optimize_prompt(prompt=prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ET3LUw5cFRQE"
   },
   "outputs": [],
   "source": [
    "display(Markdown(response.suggested_prompt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FUr06qOWxuy9"
   },
   "source": [
    "## **Part 2: The Data-Driven Optimizer**\n",
    "\n",
    "The following sections will guide you through setting up your environment, preparing your data, 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": "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."
   ]
  },
  {
   "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": "32f58bd27132"
   },
   "source": [
    "Set the optimization configuration."
   ]
  },
  {
   "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",
    "    \"eval_metrics_types\": [\"question_answering_correctness\", \"fluency\"],\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": "viiv5r23KC-X"
   },
   "source": [
    "## (Optionally) Visualize results with the interactive app\n",
    "\n",
    "The tutorial includes a helper function to launch a Gradio-based web interface. This is a great way to visually explore all the different results the optimizer generated with both prompt optimization approaches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sw0k2sv6YehS"
   },
   "source": [
    "### Build the Gradio application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gZizfa1YYOio"
   },
   "outputs": [],
   "source": [
    "def find_vapo_runs(base_path: str) -> List[str]:\n",
    "    \"\"\"Find VAPO run directories containing required files.\"\"\"\n",
    "\n",
    "    REQUIRED_FILES = {\"eval_results.json\", \"templates.json\"}\n",
    "\n",
    "    bucket_name, prefix = parse_gcs_path(base_path)\n",
    "    all_objects = list_gcs_objects(base_path)\n",
    "\n",
    "    # Group files by directory\n",
    "    directories = {}\n",
    "    for obj_path in all_objects:\n",
    "        dir_path = \"/\".join(obj_path.split(\"/\")[:-1])\n",
    "        filename = obj_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",
    "    valid_runs = []\n",
    "    for dir_path, files in directories.items():\n",
    "        if REQUIRED_FILES.issubset(files):\n",
    "            valid_runs.append(f\"gs://{bucket_name}/{dir_path}\")\n",
    "\n",
    "    return valid_runs\n",
    "\n",
    "\n",
    "# ============================================================================\n",
    "# Data Processing\n",
    "# ============================================================================\n",
    "\n",
    "\n",
    "def extract_metric_name(column_name: str) -> str:\n",
    "    \"\"\"Extract clean metric name from column name.\"\"\"\n",
    "    # Try to extract from pattern like \".metric_name/\"\n",
    "    match = re.search(r\"\\.(\\w+)/\", column_name)\n",
    "    if match:\n",
    "        return match.group(1)\n",
    "\n",
    "    # Fallback to splitting approach\n",
    "    parts = column_name.split(\".\")\n",
    "    if parts:\n",
    "        return parts[-1].split(\"/\")[0]\n",
    "\n",
    "    return column_name\n",
    "\n",
    "\n",
    "def is_metric_column(col: str) -> bool:\n",
    "    \"\"\"Check if a column is a metric column with score, confidence, or explanation.\"\"\"\n",
    "    # Must end with /score, /confidence, /explanation (or with .)\n",
    "    return bool(re.search(r\"[/\\.](score|confidence|explanation)$\", col))\n",
    "\n",
    "\n",
    "def is_core_field(col: str) -> bool:\n",
    "    \"\"\"Check if a column is a core field (not a metric).\"\"\"\n",
    "    # Exact core fields we want to keep\n",
    "    CORE_FIELDS = [\n",
    "        \"question\",\n",
    "        \"target\",\n",
    "        \"ctx\",\n",
    "        \"context\",\n",
    "        \"prompt\",\n",
    "        \"response\",\n",
    "        \"reference\",\n",
    "    ]\n",
    "\n",
    "    col_lower = col.lower()\n",
    "\n",
    "    # If it has metric indicators, it's not a core field\n",
    "    if is_metric_column(col):\n",
    "        return False\n",
    "\n",
    "    # Check if it exactly matches a core field\n",
    "    return col_lower in CORE_FIELDS\n",
    "\n",
    "\n",
    "def is_duplicate_text_field(col: str, all_columns: List[str]) -> bool:\n",
    "    \"\"\"Check if a column is a duplicate text field that should be excluded.\"\"\"\n",
    "\n",
    "    # If it's a metric column (ends with score/confidence/explanation), keep it\n",
    "    if is_metric_column(col):\n",
    "        return False\n",
    "\n",
    "    # Core fields that might have duplicates\n",
    "    CORE_FIELDS = [\n",
    "        \"question\",\n",
    "        \"target\",\n",
    "        \"ctx\",\n",
    "        \"context\",\n",
    "        \"prompt\",\n",
    "        \"response\",\n",
    "        \"reference\",\n",
    "    ]\n",
    "\n",
    "    col_lower = col.lower()\n",
    "\n",
    "    # Check if this column starts with or contains a core field name followed by underscore\n",
    "    # This catches patterns like question_fluency, question_answering_correctness, etc.\n",
    "    for core_field in CORE_FIELDS:\n",
    "        if col_lower.startswith(f\"{core_field}_\") or f\"_{core_field}\" in col_lower:\n",
    "            # It's a duplicate if it's not the exact core field\n",
    "            return col_lower != core_field\n",
    "\n",
    "    # Also check for fields that end with _fluency, _correctness, etc. (common metric names)\n",
    "    # but aren't actual metrics (don't end with /score, /confidence, /explanation)\n",
    "    metric_names = [\"fluency\", \"correctness\", \"relevance\", \"coherence\"]\n",
    "    for metric_name in metric_names:\n",
    "        if col_lower.endswith(f\"_{metric_name}\") or col_lower == metric_name:\n",
    "            # Check if there's a corresponding metric column\n",
    "            # If there are columns like fluency/score, then 'fluency' alone is likely a duplicate\n",
    "            has_metric_version = any(\n",
    "                metric_name in c.lower() and is_metric_column(c) for c in all_columns\n",
    "            )\n",
    "            if has_metric_version:\n",
    "                return True\n",
    "\n",
    "    return False\n",
    "\n",
    "\n",
    "def filter_columns_for_display(columns: List[str]) -> List[str]:\n",
    "    \"\"\"Filter columns to keep only core fields and metric scores/confidence/explanations.\"\"\"\n",
    "    filtered = []\n",
    "    seen_context = False\n",
    "\n",
    "    for col in columns:\n",
    "        col_lower = col.lower()\n",
    "\n",
    "        # Skip duplicate text fields\n",
    "        if is_duplicate_text_field(col, columns):\n",
    "            continue\n",
    "\n",
    "        # Handle core fields\n",
    "        if is_core_field(col):\n",
    "            # Skip 'context' if we already have 'ctx' to avoid duplication\n",
    "            if col_lower == \"context\":\n",
    "                if not seen_context and \"ctx\" not in [c.lower() for c in filtered]:\n",
    "                    filtered.append(col)\n",
    "                    seen_context = True\n",
    "            elif col_lower == \"ctx\":\n",
    "                # Remove context if it was added, prefer ctx\n",
    "                filtered = [c for c in filtered if c.lower() != \"context\"]\n",
    "                filtered.append(col)\n",
    "                seen_context = True\n",
    "            else:\n",
    "                filtered.append(col)\n",
    "        # Handle metric columns (must have score/confidence/explanation)\n",
    "        elif is_metric_column(col):\n",
    "            filtered.append(col)\n",
    "\n",
    "    return filtered\n",
    "\n",
    "\n",
    "def process_evaluation_result(result: Dict[str, Any]) -> pd.DataFrame:\n",
    "    \"\"\"Process evaluation result for clean display.\"\"\"\n",
    "    df = pd.read_json(io.StringIO(result[\"metrics_table\"]))\n",
    "    return df\n",
    "\n",
    "\n",
    "def categorize_columns(columns: List[str]) -> Dict[str, List[str]]:\n",
    "    \"\"\"Categorize columns by type for better organization.\"\"\"\n",
    "    categories = {\n",
    "        \"Core Fields\": [],\n",
    "        \"Metric Scores\": [],\n",
    "        \"Metric Confidence\": [],\n",
    "        \"Metric Explanations\": [],\n",
    "    }\n",
    "\n",
    "    for col in columns:\n",
    "        if is_core_field(col):\n",
    "            categories[\"Core Fields\"].append(col)\n",
    "        elif col.endswith(\"/score\") or col.endswith(\".score\"):\n",
    "            categories[\"Metric Scores\"].append(col)\n",
    "        elif col.endswith(\"/confidence\") or col.endswith(\".confidence\"):\n",
    "            categories[\"Metric Confidence\"].append(col)\n",
    "        elif col.endswith(\"/explanation\") or col.endswith(\".explanation\"):\n",
    "            categories[\"Metric Explanations\"].append(col)\n",
    "\n",
    "    # Remove empty categories\n",
    "    return {k: v for k, v in categories.items() if v}\n",
    "\n",
    "\n",
    "def get_default_columns(all_columns: List[str]) -> List[str]:\n",
    "    \"\"\"Get default columns to display - only the specified core fields.\"\"\"\n",
    "\n",
    "    # Specific fields to show by default (in order of preference)\n",
    "    DEFAULT_FIELDS = [\"question\", \"target\", \"ctx\", \"prompt\", \"response\", \"reference\"]\n",
    "\n",
    "    default_cols = []\n",
    "\n",
    "    # Add fields in the specified order if they exist\n",
    "    for field in DEFAULT_FIELDS:\n",
    "        for col in all_columns:\n",
    "            if col.lower() == field and col not in default_cols:\n",
    "                default_cols.append(col)\n",
    "                break\n",
    "\n",
    "    # If 'ctx' doesn't exist but 'context' does, use 'context' (but not both)\n",
    "    if not any(col.lower() == \"ctx\" for col in default_cols):\n",
    "        for col in all_columns:\n",
    "            if col.lower() == \"context\" and col not in default_cols:\n",
    "                default_cols.append(col)\n",
    "                break\n",
    "\n",
    "    return default_cols\n",
    "\n",
    "\n",
    "def simplify_dataframe_for_metrics(\n",
    "    df: pd.DataFrame, selected_columns: List[str]\n",
    ") -> pd.DataFrame:\n",
    "    \"\"\"Simplify DataFrame by filtering and formatting metric columns.\"\"\"\n",
    "    if df.empty or not selected_columns:\n",
    "        return df\n",
    "\n",
    "    # Filter to selected columns that exist\n",
    "    valid_columns = [col for col in selected_columns if col in df.columns]\n",
    "\n",
    "    if not valid_columns:\n",
    "        return pd.DataFrame()\n",
    "\n",
    "    simplified_df = df[valid_columns].copy()\n",
    "\n",
    "    # Round numeric columns for better display\n",
    "    for col in simplified_df.columns:\n",
    "        if simplified_df[col].dtype in [\"float64\", \"float32\"]:\n",
    "            # Round scores and confidence to 3 decimal places\n",
    "            if any(\n",
    "                x in col for x in [\"/score\", \".score\", \"/confidence\", \".confidence\"]\n",
    "            ):\n",
    "                simplified_df[col] = simplified_df[col].round(3)\n",
    "\n",
    "    # Truncate long text fields (explanations)\n",
    "    for col in simplified_df.columns:\n",
    "        if \"/explanation\" in col or \".explanation\" in col:\n",
    "            simplified_df[col] = simplified_df[col].apply(\n",
    "                lambda x: x[:200] + \"...\" if isinstance(x, str) and len(x) > 200 else x\n",
    "            )\n",
    "\n",
    "    return simplified_df\n",
    "\n",
    "\n",
    "# ============================================================================\n",
    "# Zero-Shot Optimization\n",
    "# ============================================================================\n",
    "\n",
    "\n",
    "class ZeroShotOptimizer:\n",
    "    \"\"\"Handle zero-shot prompt optimization using Vertex AI.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.client = None\n",
    "        self.project_name = None\n",
    "        self.location = None\n",
    "\n",
    "    def initialize(self, project_name: str, location: str = \"us-central1\") -> str:\n",
    "        \"\"\"Initialize Vertex AI client.\"\"\"\n",
    "        try:\n",
    "            self.project_name = project_name\n",
    "            self.location = location\n",
    "            vertexai.init(project=project_name, location=location)\n",
    "            self.client = vertexai.Client(project=project_name, location=location)\n",
    "            return f\"✅ Successfully initialized for project: {project_name} in location: {location}\"\n",
    "        except Exception as e:\n",
    "            return f\"❌ Error: {str(e)}\"\n",
    "\n",
    "    def optimize(self, prompt: str) -> Dict[str, Any]:\n",
    "        \"\"\"Run zero-shot optimization on prompt.\"\"\"\n",
    "        if not self.client:\n",
    "            return {\n",
    "                \"error\": \"Client not initialized. Please set up the project first.\",\n",
    "                \"original_prompt\": prompt,\n",
    "                \"suggested_prompt\": \"\",\n",
    "                \"guidelines\": [],\n",
    "            }\n",
    "\n",
    "        try:\n",
    "            output = self.client.prompt_optimizer.optimize_prompt(prompt=prompt)\n",
    "\n",
    "            guidelines = [\n",
    "                {\n",
    "                    \"guideline\": g.applicable_guideline,\n",
    "                    \"improvement\": g.suggested_improvement,\n",
    "                    \"before\": g.text_before_change,\n",
    "                    \"after\": g.text_after_change,\n",
    "                }\n",
    "                for g in output.applicable_guidelines\n",
    "            ]\n",
    "\n",
    "            return {\n",
    "                \"original_prompt\": output.original_prompt,\n",
    "                \"suggested_prompt\": output.suggested_prompt,\n",
    "                \"optimization_type\": output.optimization_type,\n",
    "                \"guidelines\": guidelines,\n",
    "                \"error\": None,\n",
    "            }\n",
    "        except Exception as e:\n",
    "            return {\n",
    "                \"error\": f\"Optimization failed: {str(e)}\",\n",
    "                \"original_prompt\": prompt,\n",
    "                \"suggested_prompt\": \"\",\n",
    "                \"guidelines\": [],\n",
    "            }\n",
    "\n",
    "\n",
    "# ============================================================================\n",
    "# UI Components\n",
    "# ============================================================================\n",
    "\n",
    "\n",
    "class VAPOResultsViewer:\n",
    "    \"\"\"Main application class for VAPO results viewer.\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.optimizer = ZeroShotOptimizer()\n",
    "        self.reset_state()\n",
    "\n",
    "    def reset_state(self):\n",
    "        \"\"\"Reset all application state.\"\"\"\n",
    "        self.base_path = None\n",
    "        self.runs = []\n",
    "        self.templates = []\n",
    "        self.eval_results = []\n",
    "        self.current_run = None\n",
    "        self.current_eval_full = None  # Store full evaluation data\n",
    "        self.available_columns = []  # Store available columns\n",
    "        self.filtered_columns = []  # Store filtered columns\n",
    "\n",
    "    # ========================================================================\n",
    "    # Zero-Shot Optimization Methods\n",
    "    # ========================================================================\n",
    "\n",
    "    def setup_vertex_ai(self, project_name: str, location: str) -> str:\n",
    "        \"\"\"Setup Vertex AI client.\"\"\"\n",
    "        if not project_name:\n",
    "            return \"⚠️ Please enter a project name\"\n",
    "        if not location:\n",
    "            location = \"us-central1\"  # Default location\n",
    "        return self.optimizer.initialize(project_name, location)\n",
    "\n",
    "    def optimize_prompt(self, prompt: str) -> Tuple:\n",
    "        \"\"\"Run zero-shot optimization and format results.\"\"\"\n",
    "        if not prompt:\n",
    "            return (\n",
    "                \"⚠️ Please enter a prompt to optimize\",\n",
    "                pd.DataFrame(),\n",
    "                \"\",\n",
    "                pd.DataFrame(),\n",
    "            )\n",
    "\n",
    "        results = self.optimizer.optimize(prompt)\n",
    "        return self._format_optimization_results(results)\n",
    "\n",
    "    def _format_optimization_results(\n",
    "        self, results: Dict[str, Any]\n",
    "    ) -> Tuple[str, pd.DataFrame, str, pd.DataFrame]:\n",
    "        \"\"\"Format optimization results for display.\"\"\"\n",
    "        if results.get(\"error\"):\n",
    "            return results[\"error\"], pd.DataFrame(), \"\", pd.DataFrame()\n",
    "\n",
    "        # Create summary\n",
    "        summary = self._create_optimization_summary(results)\n",
    "\n",
    "        # Format guidelines table\n",
    "        guidelines_df = self._create_guidelines_dataframe(results[\"guidelines\"])\n",
    "\n",
    "        # Format changes table\n",
    "        changes_df = self._create_changes_dataframe(results[\"guidelines\"])\n",
    "\n",
    "        return summary, guidelines_df, results[\"suggested_prompt\"], changes_df\n",
    "\n",
    "    def _create_optimization_summary(self, results: Dict[str, Any]) -> str:\n",
    "        \"\"\"Create formatted summary of optimization results.\"\"\"\n",
    "        original_len = len(results[\"original_prompt\"])\n",
    "        optimized_len = len(results[\"suggested_prompt\"])\n",
    "        change_percent = (\n",
    "            (optimized_len / original_len * 100 - 100) if original_len > 0 else 0\n",
    "        )\n",
    "\n",
    "        guidelines_list = \"\\n\".join(\n",
    "            f\"{i}. **{g['guideline']}**\" for i, g in enumerate(results[\"guidelines\"], 1)\n",
    "        )\n",
    "\n",
    "        return f\"\"\"\n",
    "## 📊 Optimization Results\n",
    "\n",
    "### Summary\n",
    "- **Type**: `{results.get('optimization_type', 'zero_shot_prompt_optimization')}`\n",
    "- **Guidelines Applied**: **{len(results['guidelines'])}**\n",
    "- **Original Length**: {original_len} characters\n",
    "- **Optimized Length**: {optimized_len} characters\n",
    "- **Change**: {optimized_len - original_len:+d} chars ({change_percent:+.1f}%)\n",
    "\n",
    "### Applied Guidelines\n",
    "{guidelines_list if guidelines_list else \"No guidelines applied.\"}\n",
    "        \"\"\"\n",
    "\n",
    "    def _create_guidelines_dataframe(self, guidelines: List[Dict]) -> pd.DataFrame:\n",
    "        \"\"\"Create DataFrame for guidelines display.\"\"\"\n",
    "        if not guidelines:\n",
    "            return pd.DataFrame()\n",
    "\n",
    "        data = [\n",
    "            {\n",
    "                \"#\": i,\n",
    "                \"Guideline\": g[\"guideline\"],\n",
    "                \"Improvement\": g[\"improvement\"],\n",
    "            }\n",
    "            for i, g in enumerate(guidelines, 1)\n",
    "        ]\n",
    "\n",
    "        return pd.DataFrame(data)\n",
    "\n",
    "    def _create_changes_dataframe(self, guidelines: List[Dict]) -> pd.DataFrame:\n",
    "        \"\"\"Create DataFrame for before/after changes.\"\"\"\n",
    "        if not guidelines:\n",
    "            return pd.DataFrame()\n",
    "\n",
    "        data = [\n",
    "            {\"#\": i, \"Before\": g[\"before\"] or \"N/A\", \"After\": g[\"after\"] or \"N/A\"}\n",
    "            for i, g in enumerate(guidelines, 1)\n",
    "        ]\n",
    "\n",
    "        return pd.DataFrame(data)\n",
    "\n",
    "    # ========================================================================\n",
    "    # Data-Driven Optimization Methods\n",
    "    # ========================================================================\n",
    "\n",
    "    def load_runs(self, base_path: str) -> gr.Dropdown:\n",
    "        \"\"\"Load available VAPO runs.\"\"\"\n",
    "        if not base_path:\n",
    "            return gr.Dropdown(choices=[], value=None)\n",
    "\n",
    "        try:\n",
    "            self.base_path = base_path\n",
    "            self.runs = find_vapo_runs(base_path)\n",
    "\n",
    "            if not self.runs:\n",
    "                return gr.Dropdown(choices=[], value=None)\n",
    "\n",
    "            return gr.Dropdown(choices=self.runs, value=self.runs[0])\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading runs: {e}\")\n",
    "            return gr.Dropdown(choices=[], value=None)\n",
    "\n",
    "    def load_run_data(\n",
    "        self, run_path: str\n",
    "    ) -> Tuple[gr.Dropdown, pd.DataFrame, pd.DataFrame, gr.CheckboxGroup]:\n",
    "        \"\"\"Load data for specific run.\"\"\"\n",
    "        if not run_path:\n",
    "            return (\n",
    "                gr.Dropdown(choices=[], value=None),\n",
    "                pd.DataFrame(),\n",
    "                pd.DataFrame(),\n",
    "                gr.CheckboxGroup(choices=[], value=[]),\n",
    "            )\n",
    "\n",
    "        try:\n",
    "            self.current_run = run_path\n",
    "\n",
    "            # Load templates and evaluation results\n",
    "            templates_data = self._load_json(f\"{run_path}/templates.json\")\n",
    "            eval_data = self._load_json(f\"{run_path}/eval_results.json\")\n",
    "\n",
    "            # Process data\n",
    "            self.templates = [pd.json_normalize(t) for t in templates_data]\n",
    "            self.eval_results = [process_evaluation_result(r) for r in eval_data]\n",
    "\n",
    "            # Handle mismatch\n",
    "            if len(self.templates) == len(self.eval_results) + 1:\n",
    "                self.templates = self.templates[1:]\n",
    "            elif len(self.templates) != len(self.eval_results):\n",
    "                raise ValueError(\n",
    "                    f\"Mismatch: {len(self.templates)} templates vs \"\n",
    "                    f\"{len(self.eval_results)} results\"\n",
    "                )\n",
    "\n",
    "            # Create template options\n",
    "            template_options = self._create_template_options()\n",
    "\n",
    "            if template_options:\n",
    "                template_df, eval_df, column_selector = self._get_template_data(0)\n",
    "                return (\n",
    "                    gr.Dropdown(choices=template_options, value=template_options[0]),\n",
    "                    template_df,\n",
    "                    eval_df,\n",
    "                    column_selector,\n",
    "                )\n",
    "\n",
    "            return (\n",
    "                gr.Dropdown(choices=[], value=None),\n",
    "                pd.DataFrame(),\n",
    "                pd.DataFrame(),\n",
    "                gr.CheckboxGroup(choices=[], value=[]),\n",
    "            )\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading run data: {e}\")\n",
    "            return (\n",
    "                gr.Dropdown(choices=[], value=None),\n",
    "                pd.DataFrame(),\n",
    "                pd.DataFrame(),\n",
    "                gr.CheckboxGroup(choices=[], value=[]),\n",
    "            )\n",
    "\n",
    "    def _load_json(self, path: str) -> Any:\n",
    "        \"\"\"Load JSON file from GCS.\"\"\"\n",
    "        with epath.Path(path).open(\"r\") as f:\n",
    "            return json.load(f)\n",
    "\n",
    "    def _create_template_options(self) -> List[str]:\n",
    "        \"\"\"Create dropdown options for templates.\"\"\"\n",
    "        options = []\n",
    "\n",
    "        for i, template_df in enumerate(self.templates):\n",
    "            metrics = []\n",
    "\n",
    "            for col in template_df.columns:\n",
    "                if \"metric\" in col and \"mean\" in col:\n",
    "                    value = template_df[col].iloc[0]\n",
    "                    metric_name = extract_metric_name(col)\n",
    "                    metrics.append(f\"{metric_name}: {value:.3f}\")\n",
    "\n",
    "            metrics_str = \" | \".join(metrics) if metrics else \"No metrics\"\n",
    "            options.append(f\"Template {i} - {metrics_str}\")\n",
    "\n",
    "        return options\n",
    "\n",
    "    def _get_template_data(\n",
    "        self, index: int\n",
    "    ) -> Tuple[pd.DataFrame, pd.DataFrame, gr.CheckboxGroup]:\n",
    "        \"\"\"Get template and evaluation data for index.\"\"\"\n",
    "        if not (0 <= index < len(self.templates)):\n",
    "            return (\n",
    "                pd.DataFrame(),\n",
    "                pd.DataFrame(),\n",
    "                gr.CheckboxGroup(choices=[], value=[]),\n",
    "            )\n",
    "\n",
    "        # Transpose template for better display\n",
    "        template_df = self.templates[index].T.reset_index()\n",
    "        template_df.columns = [\"Field\", \"Value\"]\n",
    "\n",
    "        # Store full evaluation data\n",
    "        self.current_eval_full = self.eval_results[index]\n",
    "        self.available_columns = list(self.current_eval_full.columns)\n",
    "\n",
    "        # Filter columns to only keep core fields and metric score/confidence/explanation\n",
    "        self.filtered_columns = filter_columns_for_display(self.available_columns)\n",
    "\n",
    "        # Get default columns to display (only the specified core fields)\n",
    "        default_columns = get_default_columns(self.filtered_columns)\n",
    "\n",
    "        # Create categorized column selector\n",
    "        categorized = categorize_columns(self.filtered_columns)\n",
    "\n",
    "        # Create formatted choices with categories\n",
    "        choices_info = []\n",
    "        for category, cols in categorized.items():\n",
    "            choices_info.append(f\"📁 {category}: {len(cols)} columns\")\n",
    "\n",
    "        info_text = \" | \".join(choices_info)\n",
    "\n",
    "        # Simplify evaluation data\n",
    "        eval_df = simplify_dataframe_for_metrics(\n",
    "            self.current_eval_full, default_columns\n",
    "        )\n",
    "\n",
    "        column_selector = gr.CheckboxGroup(\n",
    "            choices=self.filtered_columns,\n",
    "            value=default_columns,\n",
    "            label=\"Select Columns to Display\",\n",
    "            info=info_text,\n",
    "        )\n",
    "\n",
    "        return template_df, eval_df, column_selector\n",
    "\n",
    "    def display_template(\n",
    "        self, template_selection: str\n",
    "    ) -> Tuple[pd.DataFrame, pd.DataFrame, gr.CheckboxGroup]:\n",
    "        \"\"\"Display selected template and evaluation results.\"\"\"\n",
    "        if not template_selection:\n",
    "            return (\n",
    "                pd.DataFrame(),\n",
    "                pd.DataFrame(),\n",
    "                gr.CheckboxGroup(choices=[], value=[]),\n",
    "            )\n",
    "\n",
    "        try:\n",
    "            index = int(template_selection.split()[1])\n",
    "            return self._get_template_data(index)\n",
    "        except (ValueError, IndexError):\n",
    "            return (\n",
    "                pd.DataFrame(),\n",
    "                pd.DataFrame(),\n",
    "                gr.CheckboxGroup(choices=[], value=[]),\n",
    "            )\n",
    "\n",
    "    def update_evaluation_display(self, selected_columns: List[str]) -> pd.DataFrame:\n",
    "        \"\"\"Update evaluation display based on selected columns.\"\"\"\n",
    "        if not selected_columns or self.current_eval_full is None:\n",
    "            return pd.DataFrame()\n",
    "\n",
    "        return simplify_dataframe_for_metrics(self.current_eval_full, selected_columns)\n",
    "\n",
    "    def add_metric_scores(\n",
    "        self, current_columns: List[str]\n",
    "    ) -> Tuple[List[str], pd.DataFrame]:\n",
    "        \"\"\"Add all metric score columns to current selection.\"\"\"\n",
    "        if not self.filtered_columns:\n",
    "            return current_columns, pd.DataFrame()\n",
    "\n",
    "        # Get all metric score columns only\n",
    "        metric_columns = []\n",
    "        for col in self.filtered_columns:\n",
    "            if col.endswith(\"/score\") or col.endswith(\".score\"):\n",
    "                metric_columns.append(col)\n",
    "\n",
    "        # Combine with current selection\n",
    "        updated = list(set(current_columns + metric_columns))\n",
    "\n",
    "        # Sort to maintain consistent order\n",
    "        updated.sort(\n",
    "            key=lambda x: self.filtered_columns.index(x)\n",
    "            if x in self.filtered_columns\n",
    "            else len(self.filtered_columns)\n",
    "        )\n",
    "\n",
    "        # Update display\n",
    "        df = self.update_evaluation_display(updated)\n",
    "\n",
    "        return updated, df\n",
    "\n",
    "    def add_metric_confidence(\n",
    "        self, current_columns: List[str]\n",
    "    ) -> Tuple[List[str], pd.DataFrame]:\n",
    "        \"\"\"Add all metric confidence columns to current selection.\"\"\"\n",
    "        if not self.filtered_columns:\n",
    "            return current_columns, pd.DataFrame()\n",
    "\n",
    "        # Get all metric confidence columns\n",
    "        confidence_columns = []\n",
    "        for col in self.filtered_columns:\n",
    "            if col.endswith(\"/confidence\") or col.endswith(\".confidence\"):\n",
    "                confidence_columns.append(col)\n",
    "\n",
    "        # Combine with current selection\n",
    "        updated = list(set(current_columns + confidence_columns))\n",
    "\n",
    "        # Sort to maintain consistent order\n",
    "        updated.sort(\n",
    "            key=lambda x: self.filtered_columns.index(x)\n",
    "            if x in self.filtered_columns\n",
    "            else len(self.filtered_columns)\n",
    "        )\n",
    "\n",
    "        # Update display\n",
    "        df = self.update_evaluation_display(updated)\n",
    "\n",
    "        return updated, df\n",
    "\n",
    "    def add_metric_explanations(\n",
    "        self, current_columns: List[str]\n",
    "    ) -> Tuple[List[str], pd.DataFrame]:\n",
    "        \"\"\"Add all metric explanation columns to current selection.\"\"\"\n",
    "        if not self.filtered_columns:\n",
    "            return current_columns, pd.DataFrame()\n",
    "\n",
    "        # Get all metric explanation columns\n",
    "        explanation_columns = []\n",
    "        for col in self.filtered_columns:\n",
    "            if col.endswith(\"/explanation\") or col.endswith(\".explanation\"):\n",
    "                explanation_columns.append(col)\n",
    "\n",
    "        # Combine with current selection\n",
    "        updated = list(set(current_columns + explanation_columns))\n",
    "\n",
    "        # Sort to maintain consistent order\n",
    "        updated.sort(\n",
    "            key=lambda x: self.filtered_columns.index(x)\n",
    "            if x in self.filtered_columns\n",
    "            else len(self.filtered_columns)\n",
    "        )\n",
    "\n",
    "        # Update display\n",
    "        df = self.update_evaluation_display(updated)\n",
    "\n",
    "        return updated, df\n",
    "\n",
    "    def reset_to_default_columns(self) -> Tuple[List[str], pd.DataFrame]:\n",
    "        \"\"\"Reset column selection to defaults.\"\"\"\n",
    "        if not self.filtered_columns:\n",
    "            return [], pd.DataFrame()\n",
    "\n",
    "        default_cols = get_default_columns(self.filtered_columns)\n",
    "        df = self.update_evaluation_display(default_cols)\n",
    "\n",
    "        return default_cols, df\n",
    "\n",
    "    def select_only_scores(self) -> Tuple[List[str], pd.DataFrame]:\n",
    "        \"\"\"Select only core fields and metric scores.\"\"\"\n",
    "        if not self.filtered_columns:\n",
    "            return [], pd.DataFrame()\n",
    "\n",
    "        # Get core fields\n",
    "        core_fields = []\n",
    "        for col in self.filtered_columns:\n",
    "            if is_core_field(col):\n",
    "                core_fields.append(col)\n",
    "\n",
    "        # Get metric scores only\n",
    "        score_columns = []\n",
    "        for col in self.filtered_columns:\n",
    "            if col.endswith(\"/score\") or col.endswith(\".score\"):\n",
    "                score_columns.append(col)\n",
    "\n",
    "        selected = core_fields + score_columns\n",
    "        df = self.update_evaluation_display(selected)\n",
    "\n",
    "        return selected, df\n",
    "\n",
    "    def clear_all(self) -> Tuple:\n",
    "        \"\"\"Clear all data and reset interface.\"\"\"\n",
    "        self.reset_state()\n",
    "\n",
    "        return (\n",
    "            \"\",  # base_path_input\n",
    "            gr.Dropdown(choices=[], value=None),  # run_dropdown\n",
    "            gr.Dropdown(choices=[], value=None),  # template_dropdown\n",
    "            pd.DataFrame(),  # template_display\n",
    "            pd.DataFrame(),  # eval_display\n",
    "            gr.CheckboxGroup(choices=[], value=[]),  # column_selector\n",
    "            \"\",  # project_input\n",
    "            \"\",  # location_input\n",
    "            \"\",  # prompt_input\n",
    "            \"\",  # optimization_summary\n",
    "            pd.DataFrame(),  # guidelines_table\n",
    "            \"\",  # suggested_prompt_display\n",
    "            pd.DataFrame(),  # changes_table\n",
    "        )\n",
    "\n",
    "    # ========================================================================\n",
    "    # UI Creation\n",
    "    # ========================================================================\n",
    "\n",
    "    def create_interface(self) -> gr.Blocks:\n",
    "        \"\"\"Create the Gradio interface.\"\"\"\n",
    "        css = \"\"\"\n",
    "        .markdown-text {\n",
    "            font-size: 14px;\n",
    "            line-height: 1.6;\n",
    "        }\n",
    "        .prompt-text {\n",
    "            font-family: 'Courier New', monospace;\n",
    "            white-space: pre-wrap;\n",
    "            word-wrap: break-word;\n",
    "            background-color: #f5f5f5;\n",
    "            padding: 10px;\n",
    "            border-radius: 5px;\n",
    "        }\n",
    "        .column-selector {\n",
    "            max-height: 300px;\n",
    "            overflow-y: auto;\n",
    "            border: 1px solid #ddd;\n",
    "            padding: 10px;\n",
    "            border-radius: 5px;\n",
    "        }\n",
    "        .metric-info {\n",
    "            background-color: #e3f2fd;\n",
    "            padding: 8px;\n",
    "            border-radius: 4px;\n",
    "            margin: 5px 0;\n",
    "        }\n",
    "        \"\"\"\n",
    "\n",
    "        with gr.Blocks(\n",
    "            title=\"VAPO Results Viewer\", theme=gr.themes.Soft(), css=css\n",
    "        ) as interface:\n",
    "\n",
    "            self._create_header()\n",
    "            self._create_setup_section()\n",
    "\n",
    "            gr.Markdown(\"---\")\n",
    "\n",
    "            with gr.Tabs():\n",
    "                self._create_zero_shot_tab()\n",
    "                self._create_data_driven_tab()\n",
    "\n",
    "        return interface\n",
    "\n",
    "    def _create_header(self):\n",
    "        \"\"\"Create application header.\"\"\"\n",
    "        gr.Markdown(\"# 🚀 VAPO Results Viewer\")\n",
    "        gr.Markdown(\n",
    "            \"Analyze prompts using Zero-Shot and Data-Driven optimization \"\n",
    "            \"with Vertex AI Prompt Optimizer\"\n",
    "        )\n",
    "\n",
    "    def _create_setup_section(self):\n",
    "        \"\"\"Create Vertex AI setup section.\"\"\"\n",
    "        with gr.Accordion(\"⚙️ Vertex AI Setup\", open=False):\n",
    "            with gr.Row():\n",
    "                with gr.Column(scale=2):\n",
    "                    self.project_input = gr.Textbox(\n",
    "                        label=\"GCP Project Name\",\n",
    "                        placeholder=\"your-project-name\",\n",
    "                        info=\"Required for zero-shot optimization\",\n",
    "                    )\n",
    "                with gr.Column(scale=2):\n",
    "                    self.location_input = gr.Textbox(\n",
    "                        label=\"Location\",\n",
    "                        placeholder=\"us-central1\",\n",
    "                        value=\"us-central1\",\n",
    "                        info=\"GCP region (e.g., us-central1, europe-west4)\",\n",
    "                    )\n",
    "                with gr.Column(scale=1):\n",
    "                    setup_btn = gr.Button(\"Initialize Vertex AI\", variant=\"secondary\")\n",
    "\n",
    "            self.setup_status = gr.Textbox(label=\"Status\", interactive=False)\n",
    "\n",
    "            setup_btn.click(\n",
    "                fn=self.setup_vertex_ai,\n",
    "                inputs=[self.project_input, self.location_input],\n",
    "                outputs=[self.setup_status],\n",
    "            )\n",
    "\n",
    "    def _create_zero_shot_tab(self):\n",
    "        \"\"\"Create zero-shot optimization tab.\"\"\"\n",
    "        with gr.Tab(\"🎯 Zero-Shot Optimization\"):\n",
    "            gr.Markdown(\"### Optimize prompts using Vertex AI's zero-shot optimization\")\n",
    "\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    self.prompt_input = gr.Textbox(\n",
    "                        label=\"Original Prompt\",\n",
    "                        placeholder=(\n",
    "                            \"Enter the prompt you want to optimize...\\n\\n\"\n",
    "                            \"Example:\\n\"\n",
    "                            \"You are a professional chef. Your goal is teaching \"\n",
    "                            \"how to cook healthy cooking recipes to your apprentice.\\n\\n\"\n",
    "                            \"Given a question from your apprentice and some context, \"\n",
    "                            \"provide the correct answer to the question.\"\n",
    "                        ),\n",
    "                        lines=8,\n",
    "                        max_lines=20,\n",
    "                    )\n",
    "\n",
    "                    with gr.Row():\n",
    "                        optimize_btn = gr.Button(\n",
    "                            \"🔄 Run Zero-Shot Optimization\", variant=\"primary\", scale=2\n",
    "                        )\n",
    "                        clear_btn = gr.Button(\"Clear\", variant=\"secondary\", scale=1)\n",
    "\n",
    "            # Results section\n",
    "            with gr.Row():\n",
    "                self.optimization_summary = gr.Markdown(\n",
    "                    label=\"Optimization Summary\", elem_classes=\"markdown-text\"\n",
    "                )\n",
    "\n",
    "            with gr.Tabs():\n",
    "                with gr.Tab(\"📝 Optimized Prompt\"):\n",
    "                    self.suggested_prompt = gr.Textbox(\n",
    "                        label=\"Suggested Prompt (Zero-Shot Optimized)\",\n",
    "                        lines=15,\n",
    "                        max_lines=30,\n",
    "                        interactive=False,\n",
    "                        elem_classes=\"prompt-text\",\n",
    "                    )\n",
    "\n",
    "                with gr.Tab(\"📋 Applied Guidelines\"):\n",
    "                    self.guidelines_table = gr.DataFrame(\n",
    "                        label=\"Guidelines and Improvements\",\n",
    "                        wrap=True,\n",
    "                        interactive=False,\n",
    "                        row_count=10,\n",
    "                    )\n",
    "\n",
    "                with gr.Tab(\"🔄 Before/After Changes\"):\n",
    "                    self.changes_table = gr.DataFrame(\n",
    "                        label=\"Text Changes\", wrap=True, interactive=False, row_count=10\n",
    "                    )\n",
    "\n",
    "            # Event handlers\n",
    "            optimize_btn.click(\n",
    "                fn=self.optimize_prompt,\n",
    "                inputs=[self.prompt_input],\n",
    "                outputs=[\n",
    "                    self.optimization_summary,\n",
    "                    self.guidelines_table,\n",
    "                    self.suggested_prompt,\n",
    "                    self.changes_table,\n",
    "                ],\n",
    "            )\n",
    "\n",
    "            clear_btn.click(\n",
    "                fn=lambda: (\"\", \"\", pd.DataFrame(), \"\", pd.DataFrame()),\n",
    "                outputs=[\n",
    "                    self.prompt_input,\n",
    "                    self.optimization_summary,\n",
    "                    self.guidelines_table,\n",
    "                    self.suggested_prompt,\n",
    "                    self.changes_table,\n",
    "                ],\n",
    "            )\n",
    "\n",
    "    def _create_data_driven_tab(self):\n",
    "        \"\"\"Create data-driven optimization tab.\"\"\"\n",
    "        with gr.Tab(\"📊 Data-Driven Optimization (VAPO)\"):\n",
    "            with gr.Row():\n",
    "                with gr.Column(scale=3):\n",
    "                    self.base_path_input = gr.Textbox(\n",
    "                        label=\"GCS Base Path\",\n",
    "                        placeholder=\"gs://your-bucket/vapo-results\",\n",
    "                        info=\"Enter the base GCS path containing VAPO runs\",\n",
    "                    )\n",
    "                with gr.Column(scale=1):\n",
    "                    with gr.Row():\n",
    "                        load_btn = gr.Button(\"Load Runs\", variant=\"primary\")\n",
    "                        clear_btn = gr.Button(\"Clear All\", variant=\"secondary\")\n",
    "\n",
    "            with gr.Row():\n",
    "                self.run_dropdown = gr.Dropdown(\n",
    "                    label=\"Select Run\", choices=[], interactive=True\n",
    "                )\n",
    "                self.template_dropdown = gr.Dropdown(\n",
    "                    label=\"Select Template\", choices=[], interactive=True\n",
    "                )\n",
    "\n",
    "            with gr.Tabs():\n",
    "                with gr.Tab(\"Template Details\"):\n",
    "                    self.template_display = gr.DataFrame(\n",
    "                        label=\"Template Information\", wrap=True, interactive=False\n",
    "                    )\n",
    "\n",
    "                with gr.Tab(\"Evaluation Results\"):\n",
    "                    # Info box\n",
    "                    gr.Markdown(\n",
    "                        \"\"\"\n",
    "                        <div class=\"metric-info\">\n",
    "                        💡 <b>Tip:</b> Only score, confidence, and explanation columns are shown for all metrics.\n",
    "                        Default view shows: question, target, ctx, prompt, response, and reference.\n",
    "                        </div>\n",
    "                        \"\"\",\n",
    "                        elem_classes=\"metric-info\",\n",
    "                    )\n",
    "\n",
    "                    # Column selector section\n",
    "                    with gr.Accordion(\"📊 Column Selection\", open=True):\n",
    "                        with gr.Row():\n",
    "                            self.column_selector = gr.CheckboxGroup(\n",
    "                                choices=[],\n",
    "                                value=[],\n",
    "                                label=\"Select Columns to Display\",\n",
    "                                elem_classes=\"column-selector\",\n",
    "                            )\n",
    "\n",
    "                        with gr.Row():\n",
    "                            with gr.Column(scale=1):\n",
    "                                gr.Markdown(\"**Quick Actions:**\")\n",
    "                                with gr.Row():\n",
    "                                    scores_only_btn = gr.Button(\n",
    "                                        \"📊 Scores Only\", size=\"sm\", variant=\"secondary\"\n",
    "                                    )\n",
    "                                    add_scores_btn = gr.Button(\n",
    "                                        \"➕ Add Scores\", size=\"sm\", variant=\"secondary\"\n",
    "                                    )\n",
    "                                    add_confidence_btn = gr.Button(\n",
    "                                        \"➕ Add Confidence\",\n",
    "                                        size=\"sm\",\n",
    "                                        variant=\"secondary\",\n",
    "                                    )\n",
    "                                    add_explanations_btn = gr.Button(\n",
    "                                        \"➕ Add Explanations\",\n",
    "                                        size=\"sm\",\n",
    "                                        variant=\"secondary\",\n",
    "                                    )\n",
    "\n",
    "                            with gr.Column(scale=1):\n",
    "                                gr.Markdown(\"**Selection:**\")\n",
    "                                with gr.Row():\n",
    "                                    reset_columns_btn = gr.Button(\n",
    "                                        \"🔄 Reset\", size=\"sm\", variant=\"secondary\"\n",
    "                                    )\n",
    "                                    select_all_btn = gr.Button(\n",
    "                                        \"✅ All\", size=\"sm\", variant=\"secondary\"\n",
    "                                    )\n",
    "                                    clear_selection_btn = gr.Button(\n",
    "                                        \"❌ Clear\", size=\"sm\", variant=\"secondary\"\n",
    "                                    )\n",
    "\n",
    "                    # Evaluation display\n",
    "                    self.eval_display = gr.DataFrame(\n",
    "                        label=\"Evaluation Metrics (Filtered View)\",\n",
    "                        wrap=True,\n",
    "                        interactive=False,\n",
    "                        row_count=20,\n",
    "                    )\n",
    "\n",
    "            # Event handlers\n",
    "            load_btn.click(\n",
    "                fn=self.load_runs,\n",
    "                inputs=[self.base_path_input],\n",
    "                outputs=[self.run_dropdown],\n",
    "            )\n",
    "\n",
    "            clear_btn.click(\n",
    "                fn=self.clear_all,\n",
    "                outputs=[\n",
    "                    self.base_path_input,\n",
    "                    self.run_dropdown,\n",
    "                    self.template_dropdown,\n",
    "                    self.template_display,\n",
    "                    self.eval_display,\n",
    "                    self.column_selector,\n",
    "                    self.project_input,\n",
    "                    self.location_input,\n",
    "                    self.prompt_input,\n",
    "                    self.optimization_summary,\n",
    "                    self.guidelines_table,\n",
    "                    self.suggested_prompt,\n",
    "                    self.changes_table,\n",
    "                ],\n",
    "            )\n",
    "\n",
    "            self.run_dropdown.change(\n",
    "                fn=self.load_run_data,\n",
    "                inputs=[self.run_dropdown],\n",
    "                outputs=[\n",
    "                    self.template_dropdown,\n",
    "                    self.template_display,\n",
    "                    self.eval_display,\n",
    "                    self.column_selector,\n",
    "                ],\n",
    "            )\n",
    "\n",
    "            self.template_dropdown.change(\n",
    "                fn=self.display_template,\n",
    "                inputs=[self.template_dropdown],\n",
    "                outputs=[\n",
    "                    self.template_display,\n",
    "                    self.eval_display,\n",
    "                    self.column_selector,\n",
    "                ],\n",
    "            )\n",
    "\n",
    "            # Column selector handlers\n",
    "            self.column_selector.change(\n",
    "                fn=self.update_evaluation_display,\n",
    "                inputs=[self.column_selector],\n",
    "                outputs=[self.eval_display],\n",
    "            )\n",
    "\n",
    "            scores_only_btn.click(\n",
    "                fn=self.select_only_scores,\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "            add_scores_btn.click(\n",
    "                fn=self.add_metric_scores,\n",
    "                inputs=[self.column_selector],\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "            add_confidence_btn.click(\n",
    "                fn=self.add_metric_confidence,\n",
    "                inputs=[self.column_selector],\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "            add_explanations_btn.click(\n",
    "                fn=self.add_metric_explanations,\n",
    "                inputs=[self.column_selector],\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "            reset_columns_btn.click(\n",
    "                fn=self.reset_to_default_columns,\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "            select_all_btn.click(\n",
    "                fn=lambda: (\n",
    "                    self.filtered_columns,\n",
    "                    self.update_evaluation_display(self.filtered_columns),\n",
    "                ),\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "            clear_selection_btn.click(\n",
    "                fn=lambda: ([], pd.DataFrame()),\n",
    "                outputs=[self.column_selector, self.eval_display],\n",
    "            )\n",
    "\n",
    "\n",
    "# ============================================================================\n",
    "# Application Entry Point\n",
    "# ============================================================================\n",
    "\n",
    "\n",
    "def launch_app(\n",
    "    share: bool = False,\n",
    "    server_port: int = 7861,\n",
    "    server_name: str = \"0.0.0.0\",\n",
    "    debug: bool = False,\n",
    "):\n",
    "    \"\"\"Launch the Gradio application.\"\"\"\n",
    "    viewer = VAPOResultsViewer()\n",
    "    interface = viewer.create_interface()\n",
    "    interface.launch(\n",
    "        share=share, server_port=server_port, server_name=server_name, debug=debug\n",
    "    )\n",
    "\n",
    "    return interface"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zf4iCByVYoXz"
   },
   "source": [
    "### Run the Gradio's VAPO Results Viewer app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SS4NybYBxMSd"
   },
   "outputs": [],
   "source": [
    "interface = launch_app(share=True, server_port=7861, server_name=\"0.0.0.0\", debug=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7NIg5xKw75CZ"
   },
   "source": [
    "### Cleaning up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hCFMecKz752W"
   },
   "outputs": [],
   "source": [
    "delete_app = True\n",
    "delete_job = True\n",
    "delete_bucket = True\n",
    "\n",
    "if delete_app:\n",
    "    interface.close()\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",
    "    latest_job = custom_job_list[0]\n",
    "    latest_job.delete()\n",
    "\n",
    "if delete_bucket:\n",
    "    ! gsutil -m rm -r $BUCKET_URI"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "get_started_with_vertex_ai_prompt_optimizer.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
}
