Generative AI Google Brain

gemini-prompts-prompt-optimizer-get-started-with-vertex-ai-prompt-optimizer-custom-metric

3 min read

Get started with Vertex Prompt Optimizer - Custom metric

Source notebook

Repo path: gemini/prompts/prompt_optimizer/get_started_with_vertex_ai_prompt_optimizer_custom_metric.ipynb · Open on GitHub · advanced

Optimizes a Gemini prompt with Vertex AI Prompt Optimizer using a custom Cloud Function metric.

Summary

This notebook teaches how to run Vertex AI Prompt Optimizer with a data-driven workflow and a custom evaluation metric. It prepares a QA dataset, deploys a Cloud Function that uses Gemini 2.5 Flash as an LLM-as-judge, configures weighted built-in and custom metrics, runs client.prompt_optimizer.optimize, and retrieves the best prompt from Cloud Storage outputs.

Key code patterns

Initialize Vertex AI client

import vertexai
 
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
client = vertexai.Client(project=PROJECT_ID, location=LOCATION)

Creates the Vertex AI client used to submit the Prompt Optimizer job.

Deploy custom metric function

!gcloud functions deploy 'custom_engagement_personalization_metric' \
 --gen2 \
 --runtime="python310" \
 --source={str(build_path)} \
 --entry-point=main \
 --trigger-http \
 --region={LOCATION}

Hosts the custom evaluator so Prompt Optimizer can call it during evaluation.

Gemini autorater JSON response

autorater = GenerativeModel(
    "gemini-2.5-flash",
    generation_config=GenerationConfig(
        response_mime_type="application/json",
        response_schema=metric_response_schema,
    ),
)
response = autorater.generate_content(metric_prompt)

Uses Gemini 2.5 Flash as an LLM-as-judge with a structured JSON score and explanation.

Prompt optimization config

vapo_data_settings = {
    "target_model": "gemini-2.5-flash",
    "optimization_mode": "instruction",
    "eval_metrics_types": ["question_answering_correctness", "custom_metric"],
    "eval_metrics_weights": [0.8, 0.2],
}

Combines a built-in QA correctness metric with a custom metric for weighted optimization.

Run optimizer job

vapo_data_run_config = {
    "config_path": config_path,
    "wait_for_completion": True,
    "service_account": SERVICE_ACCOUNT,
}
result = client.prompt_optimizer.optimize(method="vapo", config=vapo_data_run_config)

Starts the backend Prompt Optimizer custom job using the uploaded configuration.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Cloud Storage, Cloud Functions, Cloud Run, Artifact Registry, IAM
  • SDKs / libraries: google-cloud-aiplatform, vertexai, google.cloud.storage, pydantic, pandas, etils, requests, functions-framework

When to use this

Use this pattern when you have evaluation data and need to optimize a Gemini prompt against both built-in and custom task metrics.

Gotchas & caveats

  • Vertex AI API must be enabled for the project.
  • A Cloud Storage bucket is required for input data, config, and optimizer results.
  • The backend job uses the default Compute Engine service account and needs Vertex AI User, Storage Object Admin, Artifact Registry Reader, Cloud Run Developer, and Cloud Run Invoker roles.
  • The custom metric function is deployed as a Gen 2 Cloud Function with Python 3.10, memory, timeout, concurrency, and min instances set in the notebook.
  • The notebook pins protobuf to 4.25.3 and installs google-cloud-aiplatform>=1.108.0.
  • For reliable prompt optimization, the notebook recommends 50-100 distinct samples.
  • If no ground-truth target is provided, the notebook says to set source_model instead of adding target responses.
  • The helper function list_gcs_objects calls parse_gcs_path, which is defined later in the notebook.

Best practices

  • Use a structured Pydantic OptimizationConfig before submitting the optimizer job.
  • Validate the deployed custom metric endpoint with a test request before running optimization.
  • Use response_mime_type and response_schema to force the autorater to return JSON.
  • Store the optimizer configuration as config.json in Cloud Storage.
  • Use weighted metrics to balance question_answering_correctness and the custom engagement metric.
  • Clean up the custom job, Cloud Function, and Cloud Storage bucket after the run.

Graph View

Backlinks