chore: GenAI SDK client - Add async replay test for Prompt Optimizer

PiperOrigin-RevId: 783037244

Author: sararob

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

@@ -18,20 +18,25 @@
 
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
+import pytest
 
 
-# If you re-record this test, you will need to update the replay file to
-# include the placeholder values for config path and service account
-def test_optimize(client):
-    """Tests the optimize request parameters method."""
-
+def _raise_for_unset_env_vars() -> None:
     if not os.environ.get("VAPO_CONFIG_PATH"):
         raise ValueError("VAPO_CONFIG_PATH environment variable is not set.")
     if not os.environ.get("VAPO_SERVICE_ACCOUNT_PROJECT_NUMBER"):
         raise ValueError(
             "VAPO_SERVICE_ACCOUNT_PROJECT_NUMBER " "environment variable is not set."
         )
 
+
+# If you re-record this test, you will need to update the replay file to
+# include the placeholder values for config path and service account
+def test_optimize(client):
+    """Tests the optimize request parameters method."""
+
+    _raise_for_unset_env_vars()
+
     config = types.PromptOptimizerVAPOConfig(
         config_path=os.environ.get("VAPO_CONFIG_PATH"),
         wait_for_completion=True,
@@ -53,3 +58,46 @@ def test_optimize(client):
     globals_for_file=globals(),
     test_method="prompt_optimizer.optimize",
 )
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_optimize_async(client):
+    _raise_for_unset_env_vars()
+
+    config = types.PromptOptimizerVAPOConfig(
+        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",
+        config=config,
+    )
+    assert isinstance(job, types.CustomJob)
+    assert job.state == types.JobState.JOB_STATE_PENDING
+
+
+@pytest.mark.asyncio
+async def test_optimize_async_with_config_wait_for_completion(client, caplog):
+    _raise_for_unset_env_vars()
+
+    config = types.PromptOptimizerVAPOConfig(
+        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",
+        wait_for_completion=True,
+    )
+    job = await client.aio.prompt_optimizer.optimize(
+        method="vapo",
+        config=config,
+    )
+    assert isinstance(job, types.CustomJob)
+    assert job.state == types.JobState.JOB_STATE_PENDING
+    assert "Ignoring wait_for_completion=True" in caplog.text

vertexai/_genai/prompt_optimizer.py

@@ -825,39 +825,47 @@ async def _get_custom_job(
         self._api_client._verify_response(return_value)
         return return_value
 
+    # Todo: b/428953357 - Add example in the README.
     async def optimize(
         self,
         method: str,
         config: types.PromptOptimizerVAPOConfigOrDict,
     ) -> types.CustomJob:
         """Call async Vertex AI Prompt Optimizer (VAPO).
 
-        # Todo: b/428953357 - Add example in the README.
-          Example usage:
-          client = vertexai.Client(project=PROJECT_NAME, location='us-central1')
-          vapo_config = vertexai.types.PromptOptimizerVAPOConfig(
-              config_path="gs://you-bucket-name/your-config.json",
-              service_account=service_account,
-              wait_for_completion=True
-          )
-          job = await client.aio.prompt_optimizer.optimize(
-              method="vapo", config=vapo_config)
-
-          Args:
-            method: The method for optimizing multiple prompts (currently only
-              vapo is supported).
-            config: PromptOptimizerVAPOConfig instance containing the
-              configuration for prompt optimization.
-
-          Returns:
-            The custom job that was created.
+        Note: The `wait_for_completion` parameter in the config will be
+        ignored when using the AsyncClient, as it is not supported.
+
+        Example usage:
+        client = vertexai.Client(project=PROJECT_NAME, location='us-central1')
+        vapo_config = vertexai.types.PromptOptimizerVAPOConfig(
+            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)
+
+        Args:
+          method: The method for optimizing multiple prompts (currently only
+            vapo is supported).
+          config: PromptOptimizerVAPOConfig 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)
 
+        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: