From cd99d7091110ea75cd0a6ee9447bfc781cfaafef Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 9 Feb 2026 08:56:54 -0800 Subject: [PATCH] feat: Deprecate prompt_optimizer.optimize and prompt_optimizer.optimize_prompt in favor of prompts.launch_optimization_job and prompts.optimize PiperOrigin-RevId: 867627537 --- .../vertexai/genai/test_prompt_optimizer.py | 100 ++++++-- vertexai/_genai/_logging_utils.py | 47 ++++ vertexai/_genai/client.py | 8 - vertexai/_genai/prompt_optimizer.py | 225 +++--------------- 4 files changed, 160 insertions(+), 220 deletions(-) create mode 100644 vertexai/_genai/_logging_utils.py diff --git a/tests/unit/vertexai/genai/test_prompt_optimizer.py b/tests/unit/vertexai/genai/test_prompt_optimizer.py index 0317c54778..2f34c7447c 100644 --- a/tests/unit/vertexai/genai/test_prompt_optimizer.py +++ b/tests/unit/vertexai/genai/test_prompt_optimizer.py @@ -19,6 +19,7 @@ import vertexai from vertexai._genai import prompt_optimizer +from vertexai._genai import prompts from vertexai._genai import types from google.genai import client import pandas as pd @@ -51,7 +52,10 @@ def test_prompt_optimizer_client(self): @mock.patch.object(client.Client, "_get_api_client") @mock.patch.object(prompt_optimizer.PromptOptimizer, "_create_custom_job_resource") - def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client): + @mock.patch.object(prompts.Prompts, "_create_custom_job_resource") + def test_prompt_optimizer_optimize( + self, mock_custom_job_prompts, mock_custom_job_prompt_optimizer, mock_client + ): """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( @@ -63,11 +67,14 @@ def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client): ), ) mock_client.assert_called_once() - mock_custom_job.assert_called_once() + mock_custom_job_prompts.assert_called_once() @mock.patch.object(client.Client, "_get_api_client") @mock.patch.object(prompt_optimizer.PromptOptimizer, "_create_custom_job_resource") - def test_prompt_optimizer_optimize_nano(self, mock_custom_job, mock_client): + @mock.patch.object(prompts.Prompts, "_create_custom_job_resource") + def test_prompt_optimizer_optimize_nano( + self, mock_custom_job_prompts, mock_custom_job_prompt_optimizer, mock_client + ): """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( @@ -79,10 +86,10 @@ def test_prompt_optimizer_optimize_nano(self, mock_custom_job, mock_client): ), ) mock_client.assert_called_once() - mock_custom_job.assert_called_once() + mock_custom_job_prompts.assert_called_once() @mock.patch.object(client.Client, "_get_api_client") - @mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt") + @mock.patch.object(prompts.Prompts, "_custom_optimize") def test_prompt_optimizer_optimize_prompt( self, mock_custom_optimize_prompt, mock_client ): @@ -92,7 +99,66 @@ def test_prompt_optimizer_optimize_prompt( mock_client.assert_called_once() mock_custom_optimize_prompt.assert_called_once() - @mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt") + +class TestPrompts: + """Unit tests for the Prompts client.""" + + def setup_method(self): + 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.prompts is not None + + @mock.patch.object(client.Client, "_get_api_client") + @mock.patch.object(prompts.Prompts, "_create_custom_job_resource") + def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client): + """Test that prompt_optimizer.optimize method creates a custom job.""" + test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + test_client.prompts.launch_optimization_job( + method=types.PromptOptimizerMethod.VAPO, + config=types.PromptOptimizerConfig( + config_path="gs://ssusie-vapo-sdk-test/config.json", + wait_for_completion=False, + service_account="test-service-account", + ), + ) + mock_client.assert_called_once() + mock_custom_job.assert_called_once() + + @mock.patch.object(client.Client, "_get_api_client") + @mock.patch.object(prompts.Prompts, "_create_custom_job_resource") + def test_prompt_optimizer_optimize_nano(self, mock_custom_job, mock_client): + """Test that prompt_optimizer.optimize method creates a custom job.""" + test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + test_client.prompts.launch_optimization_job( + method=types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO, + config=types.PromptOptimizerConfig( + config_path="gs://ssusie-vapo-sdk-test/config.json", + wait_for_completion=False, + service_account="test-service-account", + ), + ) + mock_client.assert_called_once() + mock_custom_job.assert_called_once() + + @mock.patch.object(client.Client, "_get_api_client") + @mock.patch.object(prompts.Prompts, "_custom_optimize") + def test_prompt_optimizer_optimize_prompt( + self, mock_custom_optimize_prompt, mock_client + ): + """Test that prompt_optimizer.optimize_prompt method calls optimize_prompt API.""" + test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + test_client.prompts.optimize(prompt="test_prompt") + mock_client.assert_called_once() + mock_custom_optimize_prompt.assert_called_once() + + @mock.patch.object(prompts.Prompts, "_custom_optimize") def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt): """Test that prompt_optimizer.optimize method for few shot optimizer.""" df = pd.DataFrame( @@ -107,7 +173,7 @@ def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt): optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE, examples_dataframe=df, ) - test_client.prompt_optimizer.optimize_prompt( + test_client.prompts.optimize( prompt="test_prompt", config=test_config, ) @@ -120,7 +186,7 @@ def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt): mock_kwargs["config"].examples_dataframe, test_config.examples_dataframe ) - @mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt") + @mock.patch.object(prompts.Prompts, "_custom_optimize") def test_prompt_optimizer_optimize_prompt_with_optimization_target( self, mock_custom_optimize_prompt ): @@ -129,7 +195,7 @@ def test_prompt_optimizer_optimize_prompt_with_optimization_target( config = types.OptimizeConfig( optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO, ) - test_client.prompt_optimizer.optimize_prompt( + test_client.prompts.optimize( prompt="test_prompt", config=config, ) @@ -139,17 +205,17 @@ def test_prompt_optimizer_optimize_prompt_with_optimization_target( ) @pytest.mark.asyncio - @mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt") + @mock.patch.object(prompts.AsyncPrompts, "_custom_optimize") async def test_async_prompt_optimizer_optimize_prompt( self, mock_custom_optimize_prompt ): """Test that async prompt_optimizer.optimize_prompt method calls optimize_prompt API.""" test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) - await test_client.aio.prompt_optimizer.optimize_prompt(prompt="test_prompt") + await test_client.aio.prompts.optimize(prompt="test_prompt") mock_custom_optimize_prompt.assert_called_once() @pytest.mark.asyncio - @mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt") + @mock.patch.object(prompts.AsyncPrompts, "_custom_optimize") async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target( self, mock_custom_optimize_prompt ): @@ -158,7 +224,7 @@ async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target( config = types.OptimizeConfig( optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO, ) - await test_client.aio.prompt_optimizer.optimize_prompt( + await test_client.aio.prompts.optimize( prompt="test_prompt", config=config, ) @@ -168,7 +234,7 @@ async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target( ) @pytest.mark.asyncio - @mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt") + @mock.patch.object(prompts.AsyncPrompts, "_custom_optimize") async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response( self, mock_custom_optimize_prompt ): @@ -185,7 +251,7 @@ async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response( optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE, examples_dataframe=df, ) - await test_client.aio.prompt_optimizer.optimize_prompt( + await test_client.aio.prompts.optimize( prompt="test_prompt", config=config, ) @@ -195,7 +261,7 @@ async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response( ) @pytest.mark.asyncio - @mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt") + @mock.patch.object(prompts.AsyncPrompts, "_custom_optimize") async def test_async_prompt_optimizer_optimize_prompt_few_shot_rubrics( self, mock_custom_optimize_prompt ): @@ -213,7 +279,7 @@ async def test_async_prompt_optimizer_optimize_prompt_few_shot_rubrics( optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS, examples_dataframe=df, ) - await test_client.aio.prompt_optimizer.optimize_prompt( + await test_client.aio.prompts.optimize( prompt="test_prompt", config=config, ) diff --git a/vertexai/_genai/_logging_utils.py b/vertexai/_genai/_logging_utils.py new file mode 100644 index 0000000000..cee6b23df4 --- /dev/null +++ b/vertexai/_genai/_logging_utils.py @@ -0,0 +1,47 @@ +# 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. +# + +import functools +from typing import Any, Callable +from google.genai import _common +import warnings + + +def show_deprecation_warning_once( + message: str, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator to show a deprecation warning once for a function.""" + + def decorator(func: Any) -> Any: + warning_done = False + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal warning_done + if not warning_done: + warning_done = True + warnings.warn(message, DeprecationWarning, stacklevel=2) + + # Suppress ExperimentalWarning while executing the deprecated wrapper + with warnings.catch_warnings(): + # We ignore ExperimentalWarning because the user will see it + # when they migrate to the new prompts module + warnings.simplefilter("ignore", category=_common.ExperimentalWarning) + return func(*args, **kwargs) + return func(*args, **kwargs) + + return wrapper + + return decorator diff --git a/vertexai/_genai/client.py b/vertexai/_genai/client.py index 9398108ca0..8a9a9ff135 100644 --- a/vertexai/_genai/client.py +++ b/vertexai/_genai/client.py @@ -95,10 +95,6 @@ def evals(self) -> "evals_module.AsyncEvals": return self._evals.AsyncEvals(self._api_client) # type: ignore[no-any-return] @property - @_common.experimental_warning( - "The Vertex SDK GenAI prompt optimizer module is experimental, " - "and may change in future versions." - ) def prompt_optimizer(self) -> "prompt_optimizer_module.AsyncPromptOptimizer": if self._prompt_optimizer is None: self._prompt_optimizer = importlib.import_module( @@ -258,10 +254,6 @@ def evals(self) -> "evals_module.Evals": return self._evals.Evals(self._api_client) # type: ignore[no-any-return] @property - @_common.experimental_warning( - "The Vertex SDK GenAI prompt optimizer module is experimental, and may change in future " - "versions." - ) def prompt_optimizer(self) -> "prompt_optimizer_module.PromptOptimizer": if self._prompt_optimizer is None: self._prompt_optimizer = importlib.import_module( diff --git a/vertexai/_genai/prompt_optimizer.py b/vertexai/_genai/prompt_optimizer.py index 8e08f7e323..ec62e6528c 100644 --- a/vertexai/_genai/prompt_optimizer.py +++ b/vertexai/_genai/prompt_optimizer.py @@ -15,7 +15,6 @@ # Code generated by the Google Gen AI SDK generator DO NOT EDIT. -import datetime import json import logging import time @@ -28,7 +27,9 @@ from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv +from . import _logging_utils from . import _prompt_optimizer_utils +from . import prompts from . import types @@ -405,6 +406,10 @@ def _wait_for_completion(self, job_name: str) -> types.CustomJob: logger.info(f"Job completed with state: {job.state}") return job + @_logging_utils.show_deprecation_warning_once( + "The prompt_optimizer.optimize method is deprecated. Please use" + " prompts.launch_optimization_job instead." + ) def optimize( self, method: types.PromptOptimizerMethod, @@ -420,82 +425,16 @@ def optimize( Returns: The custom job that was created. """ + prompts_module = prompts.Prompts(api_client_=self._api_client) - if isinstance(config, dict): - config = types.PromptOptimizerConfig(**config) - - if not config.config_path: - raise ValueError("Config path is required.") - - _OPTIMIZER_METHOD_TO_CONTAINER_URI = { - types.PromptOptimizerMethod.VAPO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0", - types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_android_v1_0", - } - container_uri = _OPTIMIZER_METHOD_TO_CONTAINER_URI.get(method) - if not container_uri: - raise ValueError( - 'Only "VAPO" and "OPTIMIZATION_TARGET_GEMINI_NANO" ' - "methods are currently supported." - ) - - if config.optimizer_job_display_name: - display_name = config.optimizer_job_display_name - else: - timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") - display_name = f"{method.value.lower()}-optimizer-{timestamp}" - - wait_for_completion = config.wait_for_completion - bucket = "/".join(config.config_path.split("/")[:-1]) - - region = self._api_client.location - project = self._api_client.project - container_args = { - "config": config.config_path, - } - args = ["--%s=%s" % (k, v) for k, v in container_args.items()] - worker_pool_specs = [ - types.WorkerPoolSpec( - replica_count=1, - machine_spec=types.MachineSpec(machine_type="n1-standard-4"), - container_spec=types.ContainerSpec( - image_uri=container_uri, - args=args, - ), - ) - ] - - service_account = _prompt_optimizer_utils._get_service_account(config) - - job_spec = types.CustomJobSpec( - worker_pool_specs=worker_pool_specs, - base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket), - service_account=service_account, + return prompts_module.launch_optimization_job( # type: ignore[no-any-return] + method=method, config=config ) - custom_job = types.CustomJob( - display_name=display_name, - job_spec=job_spec, - ) - - job = self._create_custom_job_resource( - custom_job=custom_job, - ) - - # Get the job resource name - job_resource_name = job.name - if not job_resource_name: - raise ValueError(f"Error creating job: {job}") - job_id = job_resource_name.split("/")[-1] - logger.info("Job created: %s", job.name) - - # Construct the dashboard URL - dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}" - logger.info("View the job status at: %s", dashboard_url) - - if wait_for_completion: - job = self._wait_for_completion(job_id) - return job - + @_logging_utils.show_deprecation_warning_once( + "The prompt_optimizer.optimize_prompt method is deprecated. Please use" + " prompts.optimize instead." + ) def optimize_prompt( self, *, @@ -534,30 +473,10 @@ def optimize_prompt( Returns: The parsed response from the API request. """ + prompts_module = prompts.Prompts(api_client_=self._api_client) - if isinstance(config, dict): - config = types.OptimizeConfig(**config) - - optimization_target: Optional[types.OptimizeTarget] = None - if config is not None: - optimization_target = config.optimization_target - - final_prompt = prompt - if ( - optimization_target - == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS - or optimization_target - == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE - ): - final_prompt = _prompt_optimizer_utils._get_few_shot_prompt(prompt, config) - - # TODO: b/435653980 - replace the custom method with a generated method. - config_for_api = config.model_copy() if config else None - return self._custom_optimize_prompt( - content=genai_types.Content( - parts=[genai_types.Part(text=final_prompt)], role="user" - ), - config=config_for_api, + return prompts_module.optimize( # type: ignore[no-any-return] + prompt=prompt, config=config ) def _custom_optimize_prompt( @@ -808,6 +727,10 @@ async def _get_custom_job( return return_value # Todo: b/428953357 - Add example in the README. + @_logging_utils.show_deprecation_warning_once( + "The prompt_optimizer.optimize method is deprecated. Please use" + " prompts.launch_optimization_job instead." + ) async def optimize( self, method: types.PromptOptimizerMethod, @@ -836,85 +759,12 @@ async def optimize( Returns: The custom job that was created. """ - if isinstance(config, dict): - config = types.PromptOptimizerConfig(**config) - - if not config.config_path: - raise ValueError("Config path is required.") - - _OPTIMIZER_METHOD_TO_CONTAINER_URI = { - types.PromptOptimizerMethod.VAPO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0", - types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_android_v1_0", - } - container_uri = _OPTIMIZER_METHOD_TO_CONTAINER_URI.get(method) - if not container_uri: - raise ValueError( - 'Only "VAPO" and "OPTIMIZATION_TARGET_GEMINI_NANO" ' - "methods are currently supported." - ) - - if config.wait_for_completion: - logger.info( - "Ignoring wait_for_completion=True since the AsyncClient does not support it." - ) - - if config.optimizer_job_display_name: - display_name = config.optimizer_job_display_name - else: - timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") - display_name = f"{method.value.lower()}-optimizer-{timestamp}" - - if not config.config_path: - raise ValueError("Config path is required.") - bucket = "/".join(config.config_path.split("/")[:-1]) - - region = self._api_client.location - project = self._api_client.project - container_args = { - "config": config.config_path, - } - args = ["--%s=%s" % (k, v) for k, v in container_args.items()] - worker_pool_specs = [ - types.WorkerPoolSpec( - replica_count=1, - machine_spec=types.MachineSpec(machine_type="n1-standard-4"), - container_spec=types.ContainerSpec( - image_uri=container_uri, - args=args, - ), - ) - ] - - service_account = _prompt_optimizer_utils._get_service_account(config) - - job_spec = types.CustomJobSpec( - worker_pool_specs=worker_pool_specs, - base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket), - service_account=service_account, - ) + prompts_module = prompts.AsyncPrompts(api_client_=self._api_client) - custom_job = types.CustomJob( - display_name=display_name, - job_spec=job_spec, + return await prompts_module.launch_optimization_job( # type: ignore[no-any-return] + method=method, config=config ) - job = await self._create_custom_job_resource( - custom_job=custom_job, - ) - - # Get the job id for the dashboard url and display to the user. - job_resource_name = job.name - if not job_resource_name: - raise ValueError(f"Error creating job: {job}") - job_id = job_resource_name.split("/")[-1] - logger.info("Job created: %s", job.name) - - # Construct the dashboard URL to show to the user. - dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}" - logger.info("View the job status at: %s", dashboard_url) - - return job - async def _custom_optimize_prompt( self, *, @@ -987,6 +837,10 @@ async def _custom_optimize_prompt( ) return final_response + @_logging_utils.show_deprecation_warning_once( + "The prompt_optimizer.optimize_prompt method is deprecated. Please use" + " prompts.optimize instead." + ) async def optimize_prompt( self, *, @@ -1021,27 +875,8 @@ async def optimize_prompt( Returns: The parsed response from the API request. """ - if isinstance(config, dict): - config = types.OptimizeConfig(**config) - - optimization_target: Optional[types.OptimizeTarget] = None - if config is not None: - optimization_target = config.optimization_target + prompts_module = prompts.AsyncPrompts(api_client_=self._api_client) - final_prompt = prompt - if ( - optimization_target - == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS - or optimization_target - == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE - ): - final_prompt = _prompt_optimizer_utils._get_few_shot_prompt(prompt, config) - - # TODO: b/435653980 - replace the custom method with a generated method. - config_for_api = config.model_copy() if config else None - return await self._custom_optimize_prompt( - content=genai_types.Content( - parts=[genai_types.Part(text=final_prompt)], role="user" - ), - config=config_for_api, + return await prompts_module.optimize( # type: ignore[no-any-return] + prompt=prompt, config=config )