feat: Adding VAPO Prompt Optimizer (PO-data) to the genai SDK.

PiperOrigin-RevId: 770884761

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_prompt_optimizer.py

@@ -0,0 +1,104 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation
+import copy
+import importlib
+from unittest import mock
+
+from google.cloud import aiplatform
+import vertexai
+from google.cloud.aiplatform import initializer as aiplatform_initializer
+from google.cloud.aiplatform.compat.services import job_service_client
+from google.cloud.aiplatform.compat.types import (
+    custom_job as gca_custom_job_compat,
+)
+from google.cloud.aiplatform.compat.types import io as gca_io_compat
+from google.cloud.aiplatform.compat.types import (
+    job_state as gca_job_state_compat,
+)
+from google.cloud.aiplatform.utils import gcs_utils
+from google.genai import client
+import pytest
+
+
+_TEST_PROJECT = "test-project"
+_TEST_LOCATION = "us-central1"
+pytestmark = pytest.mark.usefixtures("google_auth_mock")
+_TEST_PROJECT_NUMBER = "12345678"
+_TEST_PARENT = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}"
+_TEST_DISPLAY_NAME = f"{_TEST_PARENT}/customJobs/12345"
+_TEST_BASE_OUTPUT_DIR = "gs://test_bucket/test_base_output_dir"
+
+_TEST_CUSTOM_JOB_PROTO = gca_custom_job_compat.CustomJob(
+    display_name=_TEST_DISPLAY_NAME,
+    job_spec={
+        "base_output_directory": gca_io_compat.GcsDestination(
+            output_uri_prefix=_TEST_BASE_OUTPUT_DIR
+        ),
+    },
+    labels={"trained_by_vertex_ai": "true"},
+)
+
+
+@pytest.fixture
+def mock_create_custom_job():
+    with mock.patch.object(
+        job_service_client.JobServiceClient, "create_custom_job"
+    ) as create_custom_job_mock:
+        custom_job_proto = copy.deepcopy(_TEST_CUSTOM_JOB_PROTO)
+        custom_job_proto.name = _TEST_DISPLAY_NAME
+        custom_job_proto.state = gca_job_state_compat.JobState.JOB_STATE_PENDING
+        create_custom_job_mock.return_value = custom_job_proto
+        yield create_custom_job_mock
+
+
+class TestPromptOptimizer:
+    """Unit tests for the Prompt Optimizer client."""
+
+    def setup_method(self):
+        importlib.reload(aiplatform_initializer)
+        importlib.reload(aiplatform)
+        importlib.reload(vertexai)
+        vertexai.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+
+    @pytest.mark.usefixtures("google_auth_mock")
+    def test_prompt_optimizer_client(self):
+        test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
+        assert test_client is not None
+        assert test_client._api_client.vertexai
+        assert test_client._api_client.project == _TEST_PROJECT
+        assert test_client._api_client.location == _TEST_LOCATION
+
+    @mock.patch.object(client.Client, "_get_api_client")
+    @mock.patch.object(
+        gcs_utils.resource_manager_utils, "get_project_number", return_value=12345
+    )
+    def test_prompt_optimizer_optimize(
+        self, mock_get_project_number, mock_client, mock_create_custom_job
+    ):
+        """Test that prompt_optimizer.optimize method creates a custom job."""
+        test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
+        test_client.prompt_optimizer.optimize(
+            method="vapo",
+            config={
+                "config_path": "gs://ssusie-vapo-sdk-test/config.json",
+                "wait_for_completion": False,
+            },
+        )
+        mock_create_custom_job.assert_called_once()
+        mock_get_project_number.assert_called_once()

vertexai/_genai/client.py

@@ -14,7 +14,6 @@
 #
 
 import importlib
-
 from typing import Optional, Union
 
 import google.auth
@@ -24,6 +23,7 @@
 
 
 class AsyncClient:
+
     """Async Client for the GenAI SDK."""
 
     def __init__(self, api_client: client.Client):
@@ -50,6 +50,8 @@ def evals(self):
                 ) from e
         return self._evals.AsyncEvals(self._api_client)
 
+    # TODO(b/424176979): add async prompt optimizer here.
+
 
 class Client:
     """Client for the GenAI SDK.
@@ -101,6 +103,7 @@ def __init__(
             http_options=http_options,
         )
         self._evals = None
+        self._prompt_optimizer = None
 
     @property
     @_common.experimental_warning(
@@ -120,3 +123,14 @@ def evals(self):
                     "google-cloud-aiplatform[evaluation]"
                 ) from e
         return self._evals.Evals(self._api_client)
+
+    @property
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI prompt optimizer module is experimental, "
+        "and may change in future versions."
+    )
+    def prompt_optimizer(self):
+        self._prompt_optimizer = importlib.import_module(
+            ".prompt_optimizer", __package__
+        )
+        return self._prompt_optimizer.PromptOptimizer(self._api_client)

vertexai/_genai/evals.py

@@ -18,13 +18,15 @@
 import logging
 from typing import Any, Callable, Optional, Union
 from urllib.parse import urlencode
+
 from google.genai import _api_module
 from google.genai import _common
 from google.genai import types as genai_types
 from google.genai._api_client import BaseApiClient
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
 import pandas as pd
+
 from . import _evals_common
 from . import types
 
@@ -1238,9 +1240,11 @@ def evaluate(
             config = types.EvaluateMethodConfig.model_validate(config)
         if isinstance(dataset, list):
             dataset = [
-                types.EvaluationDataset.model_validate(ds_item)
-                if isinstance(ds_item, dict)
-                else ds_item
+                (
+                    types.EvaluationDataset.model_validate(ds_item)
+                    if isinstance(ds_item, dict)
+                    else ds_item
+                )
                 for ds_item in dataset
             ]
         else: