chore: GenAI SDK client - fix of return type for _wait_for_completion and few mypy errors for PromptOptimizerVAPOConfigOrDict

PiperOrigin-RevId: 785528094

Author: vertex-sdk-bot

vertexai/_genai/_prompt_optimizer_utils.py

@@ -21,22 +21,19 @@ def _get_service_account(
     config: types.PromptOptimizerVAPOConfigOrDict,
 ) -> str:
     """Get the service account from the config for the custom job."""
-    if hasattr(config, "service_account") and config.service_account:
-        if (
-            hasattr(config, "service_account_project_number")
-            and config.service_account_project_number
-        ):
-            raise ValueError(
-                "Only one of service_account or service_account_project_number "
-                "can be provided."
-            )
+    if isinstance(config, dict):
+        config = types.PromptOptimizerVAPOConfig.model_validate(config)
+
+    if config.service_account and config.service_account_project_number:
+        raise ValueError(
+            "Only one of service_account or "
+            "service_account_project_number can be provided."
+        )
+    elif config.service_account:
         return config.service_account
-    elif (
-        hasattr(config, "service_account_project_number")
-        and config.service_account_project_number
-    ):
+    elif config.service_account_project_number:
         return f"{config.service_account_project_number}-compute@developer.gserviceaccount.com"
     else:
         raise ValueError(
-            "Either service_account or service_account_project_number is required."
+            "Either service_account or service_account_project_number " "is required."
         )

vertexai/_genai/prompt_optimizer.py

@@ -540,7 +540,7 @@ def _get_custom_job(
 
     """Prompt Optimizer PO-Data."""
 
-    def _wait_for_completion(self, job_name: str) -> None:
+    def _wait_for_completion(self, job_name: str) -> types.CustomJob:
 
         JOB_COMPLETE_STATES = [
             types.JobState.JOB_STATE_SUCCEEDED,
@@ -914,6 +914,8 @@ async def optimize(
 
         # Get the job id for the dashboard url and display to the user.
         job_resource_name = job.name
+        if not job_resource_name:
+            raise ValueError(f"Error creating job: {job}")
         job_id = job_resource_name.split("/")[-1]
         logger.info("Job created: %s", job.name)
 

vertexai/_genai/types.py

@@ -1081,7 +1081,10 @@ class HttpOptions(_common.BaseModel):
     )
     extra_body: Optional[dict[str, Any]] = Field(
         default=None,
-        description="""Extra parameters to add to the request body.""",
+        description="""Extra parameters to add to the request body.
+      The structure must match the backend API's request structure.
+      - VertexAI backend API docs: https://cloud.google.com/vertex-ai/docs/reference/rest
+      - GeminiAPI backend API docs: https://ai.google.dev/api/rest""",
     )
     retry_options: Optional[HttpRetryOptions] = Field(
         default=None, description="""HTTP retry options for the request."""
@@ -1110,7 +1113,10 @@ class HttpOptionsDict(TypedDict, total=False):
     """Args passed to the async HTTP client."""
 
     extra_body: Optional[dict[str, Any]]
-    """Extra parameters to add to the request body."""
+    """Extra parameters to add to the request body.
+      The structure must match the backend API's request structure.
+      - VertexAI backend API docs: https://cloud.google.com/vertex-ai/docs/reference/rest
+      - GeminiAPI backend API docs: https://ai.google.dev/api/rest"""
 
     retry_options: Optional[HttpRetryOptionsDict]
     """HTTP retry options for the request."""