From 92d8b2a1659cc40c89a651c3a1b93ae5af5a9bf1 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 28 Oct 2025 15:52:42 -0700 Subject: [PATCH] feat: GenAI SDK client - Enabling zero-shot prompt optimization for prompts from Android API by passing optimization_target=vertexai.types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO in the config PiperOrigin-RevId: 825240223 --- ...t_optimizer_optimize_prompt_return_type.py | 14 ++++ .../vertexai/genai/test_prompt_optimizer.py | 51 +++++++++++++- vertexai/_genai/prompt_optimizer.py | 66 +++++++++++++------ vertexai/_genai/types/__init__.py | 2 + vertexai/_genai/types/common.py | 13 ++++ 5 files changed, 125 insertions(+), 21 deletions(-) diff --git a/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_prompt_return_type.py b/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_prompt_return_type.py index 8429a0072d..6e1a9170d3 100644 --- a/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_prompt_return_type.py +++ b/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_prompt_return_type.py @@ -27,6 +27,20 @@ def test_optimize_prompt(client): assert response.raw_text_response +# def test_optimize_prompt_w_optimization_target(client): +# """Tests the optimize request parameters method with optimization target.""" +# from google.genai import types as genai_types +# test_prompt = "Generate system instructions for analyzing medical articles" +# response = client.prompt_optimizer.optimize_prompt( +# prompt=test_prompt, +# config=types.OptimizeConfig( +# optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO, +# ), +# ) +# assert isinstance(response, types.OptimizeResponse) +# assert response.raw_text_response + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/tests/unit/vertexai/genai/test_prompt_optimizer.py b/tests/unit/vertexai/genai/test_prompt_optimizer.py index 2b148e2fbe..e6142109c7 100644 --- a/tests/unit/vertexai/genai/test_prompt_optimizer.py +++ b/tests/unit/vertexai/genai/test_prompt_optimizer.py @@ -69,10 +69,57 @@ def test_prompt_optimizer_optimize(self, mock_custom_job, mock_client): def test_prompt_optimizer_optimize_prompt( self, mock_custom_optimize_prompt, mock_client ): - """Test that prompt_optimizer.optimize method creates a custom job.""" + """Test that prompt_optimizer.optimize_prompt method calls optimize_prompt API.""" test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) test_client.prompt_optimizer.optimize_prompt(prompt="test_prompt") mock_client.assert_called_once() mock_custom_optimize_prompt.assert_called_once() - # TODO(b/415060797): add more tests for prompt_optimizer.optimize + @mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt") + def test_prompt_optimizer_optimize_prompt_with_optimization_target( + self, mock_custom_optimize_prompt + ): + """Test that prompt_optimizer.optimize_prompt method calls _custom_optimize_prompt with optimization_target.""" + test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + config = types.OptimizeConfig( + optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO, + ) + test_client.prompt_optimizer.optimize_prompt( + prompt="test_prompt", + config=config, + ) + mock_custom_optimize_prompt.assert_called_once_with( + content=mock.ANY, + config=config, + ) + + @pytest.mark.asyncio + @mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt") + 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") + mock_custom_optimize_prompt.assert_called_once() + + @pytest.mark.asyncio + @mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt") + async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target( + self, mock_custom_optimize_prompt + ): + """Test that async prompt_optimizer.optimize_prompt calls optimize_prompt with optimization_target.""" + test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) + config = types.OptimizeConfig( + optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO, + ) + await test_client.aio.prompt_optimizer.optimize_prompt( + prompt="test_prompt", + config=config, + ) + mock_custom_optimize_prompt.assert_called_once_with( + content=mock.ANY, + config=config, + ) + + # # TODO(b/415060797): add more tests for prompt_optimizer.optimize diff --git a/vertexai/_genai/prompt_optimizer.py b/vertexai/_genai/prompt_optimizer.py index b5f7efb354..394ca66739 100644 --- a/vertexai/_genai/prompt_optimizer.py +++ b/vertexai/_genai/prompt_optimizer.py @@ -167,6 +167,22 @@ def _GetCustomJobParameters_to_vertex( return to_object +def _OptimizeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["optimization_target"]) is not None: + setv( + parent_object, + ["optimizationTarget"], + getv(from_object, ["optimization_target"]), + ) + + return to_object + + def _OptimizeRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -176,7 +192,11 @@ def _OptimizeRequestParameters_to_vertex( setv(to_object, ["content"], getv(from_object, ["content"])) if getv(from_object, ["config"]) is not None: - setv(to_object, ["config"], getv(from_object, ["config"])) + setv( + to_object, + ["config"], + _OptimizeConfig_to_vertex(getv(from_object, ["config"]), to_object), + ) return to_object @@ -468,7 +488,10 @@ def optimize( return job def optimize_prompt( - self, *, prompt: str, config: Optional[types.OptimizeConfig] = None + self, + *, + prompt: str, + config: Optional[types.OptimizeConfig] = None, ) -> types.OptimizeResponse: """Makes an API request to _optimize_prompt and returns the parsed response. @@ -480,19 +503,21 @@ def optimize_prompt( Args: prompt: The prompt to optimize. - config: The configuration for prompt optimization. Currently, config is - not supported for a single prompt optimization. + config: Optional.The configuration for prompt optimization. To optimize + prompts from Android API provide + types.OptimizeConfig( + optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO + ) Returns: The parsed response from the API request. """ - if config is not None: - raise ValueError( - "Currently, config is not supported for a single prompt optimization." - ) prompt = genai_types.Content(parts=[genai_types.Part(text=prompt)], role="user") # TODO: b/435653980 - replace the custom method with a generated method. - return self._custom_optimize_prompt(content=prompt) + return self._custom_optimize_prompt( + content=prompt, + config=config, + ) def _custom_optimize_prompt( self, @@ -511,7 +536,6 @@ def _custom_optimize_prompt( content=content, config=config, ) - request_url_dict: Optional[dict[str, str]] if not self._api_client.vertexai: raise ValueError("This method is only supported in the Vertex AI client.") @@ -850,7 +874,6 @@ async def _custom_optimize_prompt( content=content, config=config, ) - request_url_dict: Optional[dict[str, str]] if not self._api_client.vertexai: raise ValueError("This method is only supported in the Vertex AI client.") @@ -909,7 +932,10 @@ async def _custom_optimize_prompt( return final_response async def optimize_prompt( - self, *, prompt: str, config: Optional[types.OptimizeConfig] = None + self, + *, + prompt: str, + config: Optional[types.OptimizeConfig] = None, ) -> types.OptimizeResponse: """Makes an async request to _optimize_prompt and returns an optimized prompt. @@ -920,16 +946,18 @@ async def optimize_prompt( Args: prompt: The prompt to optimize. - config: The configuration for prompt optimization. Currently, config is - not supported for a single prompt optimization. + config: Optional.The configuration for prompt optimization. To optimize + prompts from Android API provide + types.OptimizeConfig( + optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO + ) Returns: The parsed response from the API request. """ - if config is not None: - raise ValueError( - "Currently, config is not supported for a single prompt optimization." - ) prompt = genai_types.Content(parts=[genai_types.Part(text=prompt)], role="user") # TODO: b/435653980 - replace the custom method with a generated method. - return await self._custom_optimize_prompt(content=prompt) + return await self._custom_optimize_prompt( + content=prompt, + config=config, + ) diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index c09c558d6f..b54d224c12 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -584,6 +584,7 @@ from .common import OptimizeResponseEndpointDict from .common import OptimizeResponseEndpointOrDict from .common import OptimizeResponseOrDict +from .common import OptimizeTarget from .common import PairwiseChoice from .common import PairwiseMetricInput from .common import PairwiseMetricInputDict @@ -1828,6 +1829,7 @@ "RubricContentType", "EvaluationRunState", "Importance", + "OptimizeTarget", "GenerateMemoriesResponseGeneratedMemoryAction", "PromptData", "PromptDataDict", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index c6a29c9c79..151adceaf8 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -331,6 +331,13 @@ class Importance(_common.CaseInSensitiveEnum): """Low importance.""" +class OptimizeTarget(_common.CaseInSensitiveEnum): + """None""" + + OPTIMIZATION_TARGET_GEMINI_NANO = "OPTIMIZATION_TARGET_GEMINI_NANO" + """The data driven prompt optimizer designer for prompts from Android core API.""" + + class GenerateMemoriesResponseGeneratedMemoryAction(_common.CaseInSensitiveEnum): """The action to take.""" @@ -3986,6 +3993,9 @@ class OptimizeConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + optimization_target: Optional[OptimizeTarget] = Field( + default=None, description="""""" + ) class OptimizeConfigDict(TypedDict, total=False): @@ -3994,6 +4004,9 @@ class OptimizeConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptionsDict] """Used to override HTTP request options.""" + optimization_target: Optional[OptimizeTarget] + """""" + OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict]