From bcdf041747eebfd5db558830df8e8259b03045f4 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Fri, 11 Jul 2025 07:09:19 -0700 Subject: [PATCH] chore: Gen AI SDK client - update types in Prompt Optimizer PiperOrigin-RevId: 781969281 --- ...est_prompt_optimizer_optimize_job_state.py | 1 + vertexai/_genai/prompt_optimizer.py | 36 +++++++++---------- 2 files changed, 17 insertions(+), 20 deletions(-) diff --git a/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py b/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py index 09bc5fc8e6..6be77a528a 100644 --- a/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py +++ b/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py @@ -44,6 +44,7 @@ def test_optimize(client): method="vapo", config=config, ) + assert isinstance(job, types.CustomJob) assert job.state == types.JobState.JOB_STATE_SUCCEEDED diff --git a/vertexai/_genai/prompt_optimizer.py b/vertexai/_genai/prompt_optimizer.py index 2b742ea069..b072a9974d 100644 --- a/vertexai/_genai/prompt_optimizer.py +++ b/vertexai/_genai/prompt_optimizer.py @@ -618,16 +618,14 @@ def optimize( } args = ["--%s=%s" % (k, v) for k, v in container_args.items()] worker_pool_specs = [ - { - "replica_count": 1, - "container_spec": { - "image_uri": container_uri, - "args": args, - }, - "machine_spec": { - "machine_type": "n1-standard-4", - }, - } + types.WorkerPoolSpec( + replica_count=1, + machine_spec=types.MachineSpec(machine_type="n1-standard-4"), + container_spec=types.ContainerSpec( + image_uri=container_uri, + args=args, + ), + ) ] service_account = _prompt_optimizer_utils._get_service_account(config) @@ -879,16 +877,14 @@ async def optimize( } args = ["--%s=%s" % (k, v) for k, v in container_args.items()] worker_pool_specs = [ - { - "replica_count": 1, - "container_spec": { - "image_uri": container_uri, - "args": args, - }, - "machine_spec": { - "machine_type": "n1-standard-4", - }, - } + types.WorkerPoolSpec( + replica_count=1, + machine_spec=types.MachineSpec(machine_type="n1-standard-4"), + container_spec=types.ContainerSpec( + image_uri=container_uri, + args=args, + ), + ) ] service_account = _prompt_optimizer_utils._get_service_account(config)