feat!: GenAI SDK client - Enabling new data driven prompt optimization for prompts from Android API by passing

tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py

@@ -14,8 +14,8 @@
 #
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
+import logging
 import os
-
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 from google.genai import types as genai_types
@@ -38,7 +38,7 @@ def test_optimize(client):
 
     _raise_for_unset_env_vars()
 
-    config = types.PromptOptimizerVAPOConfig(
+    config = types.PromptOptimizerConfig(
         config_path=os.environ.get("VAPO_CONFIG_PATH"),
         wait_for_completion=True,
         service_account_project_number=os.environ.get(
@@ -47,7 +47,33 @@ def test_optimize(client):
         optimizer_job_display_name="optimizer_job_test",
     )
     job = client.prompt_optimizer.optimize(
-        method="vapo",
+        method=types.PromptOptimizerMethod.VAPO,
+        config=config,
+    )
+    assert isinstance(job, types.CustomJob)
+    assert job.state == genai_types.JobState.JOB_STATE_SUCCEEDED
+
+
+def test_optimize_nano(client):
+    """Tests the optimize request parameters method."""
+
+    _raise_for_unset_env_vars()
+
+    config_path = os.environ.get("VAPO_CONFIG_PATH")
+    root, ext = os.path.splitext(config_path)
+    nano_path = f"{root}_nano{ext}"
+
+    config = types.PromptOptimizerConfig(
+        config_path=nano_path,
+        wait_for_completion=True,
+        service_account_project_number=os.environ.get(
+            "VAPO_SERVICE_ACCOUNT_PROJECT_NUMBER"
+        ),
+        optimizer_job_display_name="optimizer_job_test",
+    )
+
+    job = client.prompt_optimizer.optimize(
+        method=types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO,
         config=config,
     )
     assert isinstance(job, types.CustomJob)
@@ -68,15 +94,37 @@ def test_optimize(client):
 async def test_optimize_async(client):
     _raise_for_unset_env_vars()
 
-    config = types.PromptOptimizerVAPOConfig(
+    config = types.PromptOptimizerConfig(
         config_path=os.environ.get("VAPO_CONFIG_PATH"),
         service_account_project_number=os.environ.get(
             "VAPO_SERVICE_ACCOUNT_PROJECT_NUMBER"
         ),
         optimizer_job_display_name="optimizer_job_test",
     )
     job = await client.aio.prompt_optimizer.optimize(
-        method="vapo",
+        method=types.PromptOptimizerMethod.VAPO,
+        config=config,
+    )
+    assert isinstance(job, types.CustomJob)
+    assert job.state == genai_types.JobState.JOB_STATE_PENDING
+
+
+@pytest.mark.asyncio
+async def test_optimize_nano_async(client):
+    _raise_for_unset_env_vars()
+    config_path = os.environ.get("VAPO_CONFIG_PATH")
+    root, ext = os.path.splitext(config_path)
+    nano_path = f"{root}_nano{ext}"
+
+    config = types.PromptOptimizerConfig(
+        config_path=nano_path,
+        service_account_project_number=os.environ.get(
+            "VAPO_SERVICE_ACCOUNT_PROJECT_NUMBER"
+        ),
+        optimizer_job_display_name="optimizer_job_test",
+    )
+    job = await client.aio.prompt_optimizer.optimize(
+        method=types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO,
         config=config,
     )
     assert isinstance(job, types.CustomJob)
@@ -86,8 +134,9 @@ async def test_optimize_async(client):
 @pytest.mark.asyncio
 async def test_optimize_async_with_config_wait_for_completion(client, caplog):
     _raise_for_unset_env_vars()
+    caplog.set_level(logging.INFO)
 
-    config = types.PromptOptimizerVAPOConfig(
+    config = types.PromptOptimizerConfig(
         config_path=os.environ.get("VAPO_CONFIG_PATH"),
         service_account_project_number=os.environ.get(
             "VAPO_SERVICE_ACCOUNT_PROJECT_NUMBER"
@@ -96,7 +145,7 @@ async def test_optimize_async_with_config_wait_for_completion(client, caplog):
         wait_for_completion=True,
     )
     job = await client.aio.prompt_optimizer.optimize(
-        method="vapo",
+        method=types.PromptOptimizerMethod.VAPO,
         config=config,
     )
     assert isinstance(job, types.CustomJob)

tests/unit/vertexai/genai/test_prompt_optimizer.py

@@ -54,8 +54,24 @@ 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.prompt_optimizer.optimize(
-            method="vapo",
-            config=types.PromptOptimizerVAPOConfig(
+            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(prompt_optimizer.PromptOptimizer, "_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.prompt_optimizer.optimize(
+            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",

vertexai/_genai/_prompt_optimizer_utils.py

@@ -19,11 +19,11 @@
 
 
 def _get_service_account(
-    config: types.PromptOptimizerVAPOConfigOrDict,
+    config: types.PromptOptimizerConfigOrDict,
 ) -> str:
     """Get the service account from the config for the custom job."""
     if isinstance(config, dict):
-        config = types.PromptOptimizerVAPOConfig.model_validate(config)
+        config = types.PromptOptimizerConfig.model_validate(config)
 
     if config.service_account and config.service_account_project_number:
         raise ValueError(

vertexai/_genai/prompt_optimizer.py

@@ -407,37 +407,46 @@ def _wait_for_completion(self, job_name: str) -> types.CustomJob:
 
     def optimize(
         self,
-        method: str,
-        config: types.PromptOptimizerVAPOConfigOrDict,
+        method: types.PromptOptimizerMethod,
+        config: types.PromptOptimizerConfigOrDict,
     ) -> types.CustomJob:
         """Call PO-Data optimizer.
 
         Args:
-          method: The method for optimizing multiple prompts.
-          config: PromptOptimizerVAPOConfig instance containing the
+          method: The method for optimizing multiple prompts. Supported methods:
+            VAPO, OPTIMIZATION_TARGET_GEMINI_NANO.
+          config: PromptOptimizerConfig instance containing the
               configuration for prompt optimization.
         Returns:
           The custom job that was created.
         """
 
-        if method != "vapo":
-            raise ValueError("Only vapo method is currently supported.")
-
         if isinstance(config, dict):
-            config = types.PromptOptimizerVAPOConfig(**config)
+            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/android-apo:preview_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"vapo-optimizer-{timestamp}"
+            display_name = f"{method.value.lower()}-optimizer-{timestamp}"
+
         wait_for_completion = config.wait_for_completion
-        if not config.config_path:
-            raise ValueError("Config path is required.")
         bucket = "/".join(config.config_path.split("/")[:-1])
 
-        container_uri = "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0"
-
         region = self._api_client.location
         project = self._api_client.project
         container_args = {
@@ -766,8 +775,8 @@ async def _get_custom_job(
     # Todo: b/428953357 - Add example in the README.
     async def optimize(
         self,
-        method: str,
-        config: types.PromptOptimizerVAPOConfigOrDict,
+        method: types.PromptOptimizerMethod,
+        config: types.PromptOptimizerConfigOrDict,
     ) -> types.CustomJob:
         """Call async Vertex AI Prompt Optimizer (VAPO).
 
@@ -777,26 +786,37 @@ async def optimize(
 
         Example usage:
         client = vertexai.Client(project=PROJECT_NAME, location='us-central1')
-        vapo_config = vertexai.types.PromptOptimizerVAPOConfig(
+        vapo_config = vertexai.types.PromptOptimizerConfig(
             config_path='gs://you-bucket-name/your-config.json',
             service_account=service_account,
         )
         job = await client.aio.prompt_optimizer.optimize(
-            method='vapo', config=vapo_config)
+            method=types.PromptOptimizerMethod.VAPO, config=vapo_config)
 
         Args:
-          method: The method for optimizing multiple prompts (currently only
-            vapo is supported).
-          config: PromptOptimizerVAPOConfig instance containing the
+          method: The method for optimizing multiple prompts. Supported methods:
+            VAPO, OPTIMIZATION_TARGET_GEMINI_NANO.
+          config: PromptOptimizerConfig instance containing the
             configuration for prompt optimization.
         Returns:
           The custom job that was created.
         """
-        if method != "vapo":
-            raise ValueError("Only vapo methods is currently supported.")
-
         if isinstance(config, dict):
-            config = types.PromptOptimizerVAPOConfig(**config)
+            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/android-apo:preview_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(
@@ -807,14 +827,12 @@ async def optimize(
             display_name = config.optimizer_job_display_name
         else:
             timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-            display_name = f"vapo-optimizer-{timestamp}"
+            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])
 
-        container_uri = "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0"
-
         region = self._api_client.location
         project = self._api_client.project
         container_args = {

vertexai/_genai/types/__init__.py

@@ -587,6 +587,9 @@
 from .common import OptimizeResponseEndpointDict
 from .common import OptimizeResponseEndpointOrDict
 from .common import OptimizeResponseOrDict
+from .common import OptimizerMethodPlaceholder
+from .common import OptimizerMethodPlaceholderDict
+from .common import OptimizerMethodPlaceholderOrDict
 from .common import OptimizeTarget
 from .common import PairwiseChoice
 from .common import PairwiseMetricInput
@@ -618,9 +621,10 @@
 from .common import PromptDataDict
 from .common import PromptDataOrDict
 from .common import PromptDict
-from .common import PromptOptimizerVAPOConfig
-from .common import PromptOptimizerVAPOConfigDict
-from .common import PromptOptimizerVAPOConfigOrDict
+from .common import PromptOptimizerConfig
+from .common import PromptOptimizerConfigDict
+from .common import PromptOptimizerConfigOrDict
+from .common import PromptOptimizerMethod
 from .common import PromptOrDict
 from .common import PromptRef
 from .common import PromptRefDict
@@ -1739,9 +1743,12 @@
     "UpdateDatasetConfig",
     "UpdateDatasetConfigDict",
     "UpdateDatasetConfigOrDict",
-    "PromptOptimizerVAPOConfig",
-    "PromptOptimizerVAPOConfigDict",
-    "PromptOptimizerVAPOConfigOrDict",
+    "PromptOptimizerConfig",
+    "PromptOptimizerConfigDict",
+    "PromptOptimizerConfigOrDict",
+    "OptimizerMethodPlaceholder",
+    "OptimizerMethodPlaceholderDict",
+    "OptimizerMethodPlaceholderOrDict",
     "ApplicableGuideline",
     "ApplicableGuidelineDict",
     "ApplicableGuidelineOrDict",
@@ -1837,6 +1844,7 @@
     "Importance",
     "OptimizeTarget",
     "GenerateMemoriesResponseGeneratedMemoryAction",
+    "PromptOptimizerMethod",
     "PromptData",
     "PromptDataDict",
     "PromptDataOrDict",

vertexai/_genai/types/common.py

@@ -354,6 +354,15 @@ class GenerateMemoriesResponseGeneratedMemoryAction(_common.CaseInSensitiveEnum)
     """The memory was deleted."""
 
 
+class PromptOptimizerMethod(_common.CaseInSensitiveEnum):
+    """The method for data driven prompt optimization."""
+
+    VAPO = "VAPO"
+    """The default data driven Vertex AI Prompt Optimizer."""
+    OPTIMIZATION_TARGET_GEMINI_NANO = "OPTIMIZATION_TARGET_GEMINI_NANO"
+    """The data driven prompt optimizer designer for prompts from Android core API."""
+
+
 class CreateEvaluationItemConfig(_common.BaseModel):
     """Config to create an evaluation item."""
 
@@ -12025,7 +12034,7 @@ class _UpdateDatasetParametersDict(TypedDict, total=False):
 ]
 
 
-class PromptOptimizerVAPOConfig(_common.BaseModel):
+class PromptOptimizerConfig(_common.BaseModel):
     """VAPO Prompt Optimizer Config."""
 
     config_path: Optional[str] = Field(
@@ -12050,7 +12059,7 @@ class PromptOptimizerVAPOConfig(_common.BaseModel):
     )
 
 
-class PromptOptimizerVAPOConfigDict(TypedDict, total=False):
+class PromptOptimizerConfigDict(TypedDict, total=False):
     """VAPO Prompt Optimizer Config."""
 
     config_path: Optional[str]
@@ -12069,8 +12078,26 @@ class PromptOptimizerVAPOConfigDict(TypedDict, total=False):
     """The display name of the optimization job. If not provided, a display name in the format of "vapo-optimizer-{timestamp}" will be used."""
 
 
-PromptOptimizerVAPOConfigOrDict = Union[
-    PromptOptimizerVAPOConfig, PromptOptimizerVAPOConfigDict
+PromptOptimizerConfigOrDict = Union[PromptOptimizerConfig, PromptOptimizerConfigDict]
+
+
+class OptimizerMethodPlaceholder(_common.BaseModel):
+    """Placeholder class to generate OptimizerMethod enum in common.py."""
+
+    method: Optional[PromptOptimizerMethod] = Field(
+        default=None, description="""The method for optimizing multiple prompts."""
+    )
+
+
+class OptimizerMethodPlaceholderDict(TypedDict, total=False):
+    """Placeholder class to generate OptimizerMethod enum in common.py."""
+
+    method: Optional[PromptOptimizerMethod]
+    """The method for optimizing multiple prompts."""
+
+
+OptimizerMethodPlaceholderOrDict = Union[
+    OptimizerMethodPlaceholder, OptimizerMethodPlaceholderDict
 ]