Skip to content

Commit cd99d70

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
feat: Deprecate prompt_optimizer.optimize and prompt_optimizer.optimize_prompt in favor of prompts.launch_optimization_job and prompts.optimize
PiperOrigin-RevId: 867627537
1 parent 6907f89 commit cd99d70

File tree

4 files changed

+160
-220
lines changed

4 files changed

+160
-220
lines changed

‎tests/unit/vertexai/genai/test_prompt_optimizer.py‎

Lines changed: 83 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import vertexai
2121
from vertexai._genai import prompt_optimizer
22+
from vertexai._genai import prompts
2223
from vertexai._genai import types
2324
from google.genai import client
2425
import pandas as pd
@@ -51,7 +52,10 @@ def test_prompt_optimizer_client(self):
5152

5253
@mock.patch.object(client.Client, "_get_api_client")
5354
@mock.patch.object(prompt_optimizer.PromptOptimizer, "_create_custom_job_resource")
54-
def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client):
55+
@mock.patch.object(prompts.Prompts, "_create_custom_job_resource")
56+
def test_prompt_optimizer_optimize(
57+
self, mock_custom_job_prompts, mock_custom_job_prompt_optimizer, mock_client
58+
):
5559
"""Test that prompt_optimizer.optimize method creates a custom job."""
5660
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
5761
test_client.prompt_optimizer.optimize(
@@ -63,11 +67,14 @@ def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client):
6367
),
6468
)
6569
mock_client.assert_called_once()
66-
mock_custom_job.assert_called_once()
70+
mock_custom_job_prompts.assert_called_once()
6771

6872
@mock.patch.object(client.Client, "_get_api_client")
6973
@mock.patch.object(prompt_optimizer.PromptOptimizer, "_create_custom_job_resource")
70-
def test_prompt_optimizer_optimize_nano(self, mock_custom_job, mock_client):
74+
@mock.patch.object(prompts.Prompts, "_create_custom_job_resource")
75+
def test_prompt_optimizer_optimize_nano(
76+
self, mock_custom_job_prompts, mock_custom_job_prompt_optimizer, mock_client
77+
):
7178
"""Test that prompt_optimizer.optimize method creates a custom job."""
7279
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
7380
test_client.prompt_optimizer.optimize(
@@ -79,10 +86,10 @@ def test_prompt_optimizer_optimize_nano(self, mock_custom_job, mock_client):
7986
),
8087
)
8188
mock_client.assert_called_once()
82-
mock_custom_job.assert_called_once()
89+
mock_custom_job_prompts.assert_called_once()
8390

8491
@mock.patch.object(client.Client, "_get_api_client")
85-
@mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt")
92+
@mock.patch.object(prompts.Prompts, "_custom_optimize")
8693
def test_prompt_optimizer_optimize_prompt(
8794
self, mock_custom_optimize_prompt, mock_client
8895
):
@@ -92,7 +99,66 @@ def test_prompt_optimizer_optimize_prompt(
9299
mock_client.assert_called_once()
93100
mock_custom_optimize_prompt.assert_called_once()
94101

95-
@mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt")
102+
103+
class TestPrompts:
104+
"""Unit tests for the Prompts client."""
105+
106+
def setup_method(self):
107+
importlib.reload(vertexai)
108+
vertexai.init(
109+
project=_TEST_PROJECT,
110+
location=_TEST_LOCATION,
111+
)
112+
113+
@pytest.mark.usefixtures("google_auth_mock")
114+
def test_prompt_optimizer_client(self):
115+
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
116+
assert test_client.prompts is not None
117+
118+
@mock.patch.object(client.Client, "_get_api_client")
119+
@mock.patch.object(prompts.Prompts, "_create_custom_job_resource")
120+
def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client):
121+
"""Test that prompt_optimizer.optimize method creates a custom job."""
122+
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
123+
test_client.prompts.launch_optimization_job(
124+
method=types.PromptOptimizerMethod.VAPO,
125+
config=types.PromptOptimizerConfig(
126+
config_path="gs://ssusie-vapo-sdk-test/config.json",
127+
wait_for_completion=False,
128+
service_account="test-service-account",
129+
),
130+
)
131+
mock_client.assert_called_once()
132+
mock_custom_job.assert_called_once()
133+
134+
@mock.patch.object(client.Client, "_get_api_client")
135+
@mock.patch.object(prompts.Prompts, "_create_custom_job_resource")
136+
def test_prompt_optimizer_optimize_nano(self, mock_custom_job, mock_client):
137+
"""Test that prompt_optimizer.optimize method creates a custom job."""
138+
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
139+
test_client.prompts.launch_optimization_job(
140+
method=types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO,
141+
config=types.PromptOptimizerConfig(
142+
config_path="gs://ssusie-vapo-sdk-test/config.json",
143+
wait_for_completion=False,
144+
service_account="test-service-account",
145+
),
146+
)
147+
mock_client.assert_called_once()
148+
mock_custom_job.assert_called_once()
149+
150+
@mock.patch.object(client.Client, "_get_api_client")
151+
@mock.patch.object(prompts.Prompts, "_custom_optimize")
152+
def test_prompt_optimizer_optimize_prompt(
153+
self, mock_custom_optimize_prompt, mock_client
154+
):
155+
"""Test that prompt_optimizer.optimize_prompt method calls optimize_prompt API."""
156+
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
157+
test_client.prompts.optimize(prompt="test_prompt")
158+
mock_client.assert_called_once()
159+
mock_custom_optimize_prompt.assert_called_once()
160+
161+
@mock.patch.object(prompts.Prompts, "_custom_optimize")
96162
def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt):
97163
"""Test that prompt_optimizer.optimize method for few shot optimizer."""
98164
df = pd.DataFrame(
@@ -107,7 +173,7 @@ def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt):
107173
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
108174
examples_dataframe=df,
109175
)
110-
test_client.prompt_optimizer.optimize_prompt(
176+
test_client.prompts.optimize(
111177
prompt="test_prompt",
112178
config=test_config,
113179
)
@@ -120,7 +186,7 @@ def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt):
120186
mock_kwargs["config"].examples_dataframe, test_config.examples_dataframe
121187
)
122188

123-
@mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt")
189+
@mock.patch.object(prompts.Prompts, "_custom_optimize")
124190
def test_prompt_optimizer_optimize_prompt_with_optimization_target(
125191
self, mock_custom_optimize_prompt
126192
):
@@ -129,7 +195,7 @@ def test_prompt_optimizer_optimize_prompt_with_optimization_target(
129195
config = types.OptimizeConfig(
130196
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO,
131197
)
132-
test_client.prompt_optimizer.optimize_prompt(
198+
test_client.prompts.optimize(
133199
prompt="test_prompt",
134200
config=config,
135201
)
@@ -139,17 +205,17 @@ def test_prompt_optimizer_optimize_prompt_with_optimization_target(
139205
)
140206

141207
@pytest.mark.asyncio
142-
@mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt")
208+
@mock.patch.object(prompts.AsyncPrompts, "_custom_optimize")
143209
async def test_async_prompt_optimizer_optimize_prompt(
144210
self, mock_custom_optimize_prompt
145211
):
146212
"""Test that async prompt_optimizer.optimize_prompt method calls optimize_prompt API."""
147213
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
148-
await test_client.aio.prompt_optimizer.optimize_prompt(prompt="test_prompt")
214+
await test_client.aio.prompts.optimize(prompt="test_prompt")
149215
mock_custom_optimize_prompt.assert_called_once()
150216

151217
@pytest.mark.asyncio
152-
@mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt")
218+
@mock.patch.object(prompts.AsyncPrompts, "_custom_optimize")
153219
async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target(
154220
self, mock_custom_optimize_prompt
155221
):
@@ -158,7 +224,7 @@ async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target(
158224
config = types.OptimizeConfig(
159225
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO,
160226
)
161-
await test_client.aio.prompt_optimizer.optimize_prompt(
227+
await test_client.aio.prompts.optimize(
162228
prompt="test_prompt",
163229
config=config,
164230
)
@@ -168,7 +234,7 @@ async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target(
168234
)
169235

170236
@pytest.mark.asyncio
171-
@mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt")
237+
@mock.patch.object(prompts.AsyncPrompts, "_custom_optimize")
172238
async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response(
173239
self, mock_custom_optimize_prompt
174240
):
@@ -185,7 +251,7 @@ async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response(
185251
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
186252
examples_dataframe=df,
187253
)
188-
await test_client.aio.prompt_optimizer.optimize_prompt(
254+
await test_client.aio.prompts.optimize(
189255
prompt="test_prompt",
190256
config=config,
191257
)
@@ -195,7 +261,7 @@ async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response(
195261
)
196262

197263
@pytest.mark.asyncio
198-
@mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt")
264+
@mock.patch.object(prompts.AsyncPrompts, "_custom_optimize")
199265
async def test_async_prompt_optimizer_optimize_prompt_few_shot_rubrics(
200266
self, mock_custom_optimize_prompt
201267
):
@@ -213,7 +279,7 @@ async def test_async_prompt_optimizer_optimize_prompt_few_shot_rubrics(
213279
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS,
214280
examples_dataframe=df,
215281
)
216-
await test_client.aio.prompt_optimizer.optimize_prompt(
282+
await test_client.aio.prompts.optimize(
217283
prompt="test_prompt",
218284
config=config,
219285
)

‎vertexai/_genai/_logging_utils.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
16+
import functools
17+
from typing import Any, Callable
18+
from google.genai import _common
19+
import warnings
20+
21+
22+
def show_deprecation_warning_once(
23+
message: str,
24+
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
25+
"""Decorator to show a deprecation warning once for a function."""
26+
27+
def decorator(func: Any) -> Any:
28+
warning_done = False
29+
30+
@functools.wraps(func)
31+
def wrapper(*args: Any, **kwargs: Any) -> Any:
32+
nonlocal warning_done
33+
if not warning_done:
34+
warning_done = True
35+
warnings.warn(message, DeprecationWarning, stacklevel=2)
36+
37+
# Suppress ExperimentalWarning while executing the deprecated wrapper
38+
with warnings.catch_warnings():
39+
# We ignore ExperimentalWarning because the user will see it
40+
# when they migrate to the new prompts module
41+
warnings.simplefilter("ignore", category=_common.ExperimentalWarning)
42+
return func(*args, **kwargs)
43+
return func(*args, **kwargs)
44+
45+
return wrapper
46+
47+
return decorator

‎vertexai/_genai/client.py‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -95,10 +95,6 @@ def evals(self) -> "evals_module.AsyncEvals":
9595
return self._evals.AsyncEvals(self._api_client) # type: ignore[no-any-return]
9696

9797
@property
98-
@_common.experimental_warning(
99-
"The Vertex SDK GenAI prompt optimizer module is experimental, "
100-
"and may change in future versions."
101-
)
10298
def prompt_optimizer(self) -> "prompt_optimizer_module.AsyncPromptOptimizer":
10399
if self._prompt_optimizer is None:
104100
self._prompt_optimizer = importlib.import_module(
@@ -258,10 +254,6 @@ def evals(self) -> "evals_module.Evals":
258254
return self._evals.Evals(self._api_client) # type: ignore[no-any-return]
259255

260256
@property
261-
@_common.experimental_warning(
262-
"The Vertex SDK GenAI prompt optimizer module is experimental, and may change in future "
263-
"versions."
264-
)
265257
def prompt_optimizer(self) -> "prompt_optimizer_module.PromptOptimizer":
266258
if self._prompt_optimizer is None:
267259
self._prompt_optimizer = importlib.import_module(

0 commit comments

Comments
 (0)