feat: GenAI SDK client - Enabling zero-shot prompt optimization for prompts from Android API by passing gemini_nano_config = vertexai.types.OptimizerConfig(optimization_target=vertexai.types.OptimizeTarget.GEMINI_NANO) in the config

PiperOrigin-RevId: 821852727

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_prompt_optimizer.py

@@ -69,10 +69,28 @@ 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 method calls _custom_optimize_prompt."""
         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()
 
+    @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.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,
+        )
+
     # TODO(b/415060797): add more tests for prompt_optimizer.optimize

vertexai/_genai/_prompt_optimizer_utils.py

@@ -15,9 +15,40 @@
 """Utility functions for prompt optimizer."""
 
 import json
+from google.genai import types as genai_types
 from . import types
 
 
+def _get_parameter_model(
+        config: types.OptimizeConfigOrDict | None,
+        content: genai_types.ContentOrDict | None
+        ):
+    """Get the parameters from the config and content."""
+
+    if isinstance(config, types.OptimizeConfig):
+        optimization_target = config.optimization_target
+    elif isinstance(config, dict):
+        optimization_target = config.get("optimization_target")
+    else:
+        optimization_target = None
+
+    if optimization_target is None:
+        return  types._OptimizeRequestParameters(
+                content=content,
+                config=config,
+            )
+
+    if optimization_target != types.OptimizeTarget.GEMINI_NANO:
+        raise ValueError(
+            "Currently, only `gemini_nano` optimization_target is supported."
+        )
+    return types._OptimizeRequestParameters(
+        content=content,
+        optimization_target="OPTIMIZATION_TARGET_GEMINI_NANO",
+        config=config,
+    )
+
+
 def _get_service_account(
     config: types.PromptOptimizerVAPOConfigOrDict,
 ) -> str:

vertexai/_genai/prompt_optimizer.py

@@ -178,6 +178,13 @@ def _OptimizeRequestParameters_to_vertex(
     if getv(from_object, ["config"]) is not None:
         setv(to_object, ["config"], getv(from_object, ["config"]))
 
+    if getv(from_object, ["optimization_target"]) is not None:
+        setv(
+            to_object,
+            ["optimizationTarget"],
+            getv(from_object, ["optimization_target"]),
+        )
+
     return to_object
 
 
@@ -189,6 +196,7 @@ def _optimize_prompt(
         *,
         content: Optional[genai_types.ContentOrDict] = None,
         config: Optional[types.OptimizeConfigOrDict] = None,
+        optimization_target: Optional[str] = None,
     ) -> types.OptimizeResponseEndpoint:
         """
         Optimize a single prompt.
@@ -197,6 +205,7 @@ def _optimize_prompt(
         parameter_model = types._OptimizeRequestParameters(
             content=content,
             config=config,
+            optimization_target=optimization_target,
         )
 
         request_url_dict: Optional[dict[str, str]]
@@ -468,7 +477,10 @@ def optimize(
         return job
 
     def optimize_prompt(
-        self, *, prompt: str, config: Optional[types.OptimizeConfig] = None
+        self,
+        *,
+        prompt: str,
+        config: types.OptimizeConfig | None = None,
     ) -> types.OptimizeResponse:
         """Makes an API request to _optimize_prompt and returns the parsed response.
 
@@ -481,24 +493,26 @@ 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.
+            either None or
+            types.OptimizeConfig(
+                optimization_target=types.OptimizeTarget.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,
         *,
-        content: Optional[genai_types.ContentOrDict] = None,
-        config: Optional[types.OptimizeConfigOrDict] = None,
+        content: genai_types.ContentOrDict | None,
+        config: types.OptimizeConfigOrDict | None,
     ) -> types.OptimizeResponse:
         """Optimize a single prompt.
 
@@ -507,11 +521,7 @@ def _custom_optimize_prompt(
         the parsed response.
         """
 
-        parameter_model = types._OptimizeRequestParameters(
-            content=content,
-            config=config,
-        )
-
+        parameter_model = _prompt_optimizer_utils._get_parameter_model(config, content)
         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.")
@@ -576,6 +586,7 @@ async def _optimize_prompt(
         *,
         content: Optional[genai_types.ContentOrDict] = None,
         config: Optional[types.OptimizeConfigOrDict] = None,
+        optimization_target: Optional[str] = None,
     ) -> types.OptimizeResponseEndpoint:
         """
         Optimize a single prompt.
@@ -584,6 +595,7 @@ async def _optimize_prompt(
         parameter_model = types._OptimizeRequestParameters(
             content=content,
             config=config,
+            optimization_target=optimization_target,
         )
 
         request_url_dict: Optional[dict[str, str]]
@@ -841,16 +853,12 @@ async def optimize(
     async def _custom_optimize_prompt(
         self,
         *,
-        content: Optional[genai_types.ContentOrDict] = None,
-        config: Optional[types.OptimizeConfigOrDict] = None,
+        content: genai_types.ContentOrDict | None,
+        config: types.OptimizeConfigOrDict | None,
     ) -> types.OptimizeResponse:
         """Optimize a single prompt."""
 
-        parameter_model = types._OptimizeRequestParameters(
-            content=content,
-            config=config,
-        )
-
+        parameter_model = _prompt_optimizer_utils._get_parameter_model(config, content)
         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 +917,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: types.OptimizeConfig | None = None,
     ) -> types.OptimizeResponse:
         """Makes an async request to _optimize_prompt and returns an optimized prompt.
 
@@ -921,15 +932,17 @@ 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.
+            either None or
+            types.OptimizeConfig(
+                optimization_target=types.OptimizeTarget.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,
+        )

vertexai/_genai/types/__init__.py

@@ -571,6 +571,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
@@ -1786,6 +1787,7 @@
     "RubricContentType",
     "EvaluationRunState",
     "Importance",
+    "OptimizeTarget",
     "GenerateMemoriesResponseGeneratedMemoryAction",
     "PromptData",
     "PromptDataDict",

vertexai/_genai/types/common.py

@@ -343,6 +343,13 @@ class Importance(_common.CaseInSensitiveEnum):
     """Low importance."""
 
 
+class OptimizeTarget(_common.CaseInSensitiveEnum):
+    """None"""
+
+    GEMINI_NANO = "GEMINI_NANO"
+    """The data driven prompt optimizer designer for prompts from Android core API."""
+
+
 class GenerateMemoriesResponseGeneratedMemoryAction(_common.CaseInSensitiveEnum):
     """The action to take."""
 
@@ -3998,6 +4005,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):
@@ -4006,6 +4016,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]
 
@@ -4015,6 +4028,7 @@ class _OptimizeRequestParameters(_common.BaseModel):
 
     content: Optional[genai_types.Content] = Field(default=None, description="""""")
     config: Optional[OptimizeConfig] = Field(default=None, description="""""")
+    optimization_target: Optional[str] = Field(default=None, description="""""")
 
 
 class _OptimizeRequestParametersDict(TypedDict, total=False):
@@ -4026,6 +4040,9 @@ class _OptimizeRequestParametersDict(TypedDict, total=False):
     config: Optional[OptimizeConfigDict]
     """"""
 
+    optimization_target: Optional[str]
+    """"""
+
 
 _OptimizeRequestParametersOrDict = Union[
     _OptimizeRequestParameters, _OptimizeRequestParametersDict