Skip to content

Commit 62a28d4

Browse files
sararobcopybara-github
sararob
authored and
copybara-github
committed
feat: Deprecate prompt_optimizer.optimize and prompt_optimizer.optimize_prompt in favor of prompts.launch_optimization_job and prompts.optimize
PiperOrigin-RevId: 867627537
1 parent d3924ab commit 62a28d4

File tree

1 file changed

+33
-196
lines changed

1 file changed

+33
-196
lines changed

‎vertexai/_genai/prompt_optimizer.py‎

Lines changed: 33 additions & 196 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,12 @@
1515

1616
# Code generated by the Google Gen AI SDK generator DO NOT EDIT.
1717

18-
import datetime
1918
import json
2019
import logging
2120
import time
2221
from typing import Any, Optional, Union
2322
from urllib.parse import urlencode
23+
import warnings
2424

2525
from google.genai import _api_module
2626
from google.genai import _common
@@ -29,6 +29,7 @@
2929
from google.genai._common import set_value_by_path as setv
3030

3131
from . import _prompt_optimizer_utils
32+
from . import prompts
3233
from . import types
3334

3435

@@ -420,81 +421,15 @@ def optimize(
420421
Returns:
421422
The custom job that was created.
422423
"""
423-
424-
if isinstance(config, dict):
425-
config = types.PromptOptimizerConfig(**config)
426-
427-
if not config.config_path:
428-
raise ValueError("Config path is required.")
429-
430-
_OPTIMIZER_METHOD_TO_CONTAINER_URI = {
431-
types.PromptOptimizerMethod.VAPO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0",
432-
types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_android_v1_0",
433-
}
434-
container_uri = _OPTIMIZER_METHOD_TO_CONTAINER_URI.get(method)
435-
if not container_uri:
436-
raise ValueError(
437-
'Only "VAPO" and "OPTIMIZATION_TARGET_GEMINI_NANO" '
438-
"methods are currently supported."
439-
)
440-
441-
if config.optimizer_job_display_name:
442-
display_name = config.optimizer_job_display_name
443-
else:
444-
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
445-
display_name = f"{method.value.lower()}-optimizer-{timestamp}"
446-
447-
wait_for_completion = config.wait_for_completion
448-
bucket = "/".join(config.config_path.split("/")[:-1])
449-
450-
region = self._api_client.location
451-
project = self._api_client.project
452-
container_args = {
453-
"config": config.config_path,
454-
}
455-
args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
456-
worker_pool_specs = [
457-
types.WorkerPoolSpec(
458-
replica_count=1,
459-
machine_spec=types.MachineSpec(machine_type="n1-standard-4"),
460-
container_spec=types.ContainerSpec(
461-
image_uri=container_uri,
462-
args=args,
463-
),
464-
)
465-
]
466-
467-
service_account = _prompt_optimizer_utils._get_service_account(config)
468-
469-
job_spec = types.CustomJobSpec(
470-
worker_pool_specs=worker_pool_specs,
471-
base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket),
472-
service_account=service_account,
424+
warnings.warn(
425+
"The prompt_optimizer.optimize method is deprecated. Please use"
426+
" prompts.launch_optimization_job instead.",
427+
DeprecationWarning,
428+
stacklevel=2,
473429
)
430+
prompts_module = prompts.Prompts(api_client_=self._api_client)
474431

475-
custom_job = types.CustomJob(
476-
display_name=display_name,
477-
job_spec=job_spec,
478-
)
479-
480-
job = self._create_custom_job_resource(
481-
custom_job=custom_job,
482-
)
483-
484-
# Get the job resource name
485-
job_resource_name = job.name
486-
if not job_resource_name:
487-
raise ValueError(f"Error creating job: {job}")
488-
job_id = job_resource_name.split("/")[-1]
489-
logger.info("Job created: %s", job.name)
490-
491-
# Construct the dashboard URL
492-
dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}"
493-
logger.info("View the job status at: %s", dashboard_url)
494-
495-
if wait_for_completion:
496-
job = self._wait_for_completion(job_id)
497-
return job
432+
return prompts_module.launch_optimization_job(method=method, config=config)
498433

499434
def optimize_prompt(
500435
self,
@@ -534,31 +469,15 @@ def optimize_prompt(
534469
Returns:
535470
The parsed response from the API request.
536471
"""
537-
538-
if isinstance(config, dict):
539-
config = types.OptimizeConfig(**config)
540-
541-
optimization_target: Optional[types.OptimizeTarget] = None
542-
if config is not None:
543-
optimization_target = config.optimization_target
544-
545-
final_prompt = prompt
546-
if (
547-
optimization_target
548-
== types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
549-
or optimization_target
550-
== types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
551-
):
552-
final_prompt = _prompt_optimizer_utils._get_few_shot_prompt(prompt, config)
553-
554-
# TODO: b/435653980 - replace the custom method with a generated method.
555-
config_for_api = config.model_copy() if config else None
556-
return self._custom_optimize_prompt(
557-
content=genai_types.Content(
558-
parts=[genai_types.Part(text=final_prompt)], role="user"
559-
),
560-
config=config_for_api,
472+
warnings.warn(
473+
"The prompt_optimizer.optimize_prompt method is deprecated. Please use"
474+
" prompts.optimize instead.",
475+
DeprecationWarning,
476+
stacklevel=2,
561477
)
478+
prompts_module = prompts.Prompts(api_client_=self._api_client)
479+
480+
return prompts_module.optimize(prompt=prompt, config=config)
562481

563482
def _custom_optimize_prompt(
564483
self,
@@ -836,85 +755,18 @@ async def optimize(
836755
Returns:
837756
The custom job that was created.
838757
"""
839-
if isinstance(config, dict):
840-
config = types.PromptOptimizerConfig(**config)
841-
842-
if not config.config_path:
843-
raise ValueError("Config path is required.")
844-
845-
_OPTIMIZER_METHOD_TO_CONTAINER_URI = {
846-
types.PromptOptimizerMethod.VAPO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_v1_0",
847-
types.PromptOptimizerMethod.OPTIMIZATION_TARGET_GEMINI_NANO: "us-docker.pkg.dev/vertex-ai/cair/vaipo:preview_android_v1_0",
848-
}
849-
container_uri = _OPTIMIZER_METHOD_TO_CONTAINER_URI.get(method)
850-
if not container_uri:
851-
raise ValueError(
852-
'Only "VAPO" and "OPTIMIZATION_TARGET_GEMINI_NANO" '
853-
"methods are currently supported."
854-
)
855-
856-
if config.wait_for_completion:
857-
logger.info(
858-
"Ignoring wait_for_completion=True since the AsyncClient does not support it."
859-
)
860-
861-
if config.optimizer_job_display_name:
862-
display_name = config.optimizer_job_display_name
863-
else:
864-
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
865-
display_name = f"{method.value.lower()}-optimizer-{timestamp}"
866-
867-
if not config.config_path:
868-
raise ValueError("Config path is required.")
869-
bucket = "/".join(config.config_path.split("/")[:-1])
870-
871-
region = self._api_client.location
872-
project = self._api_client.project
873-
container_args = {
874-
"config": config.config_path,
875-
}
876-
args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
877-
worker_pool_specs = [
878-
types.WorkerPoolSpec(
879-
replica_count=1,
880-
machine_spec=types.MachineSpec(machine_type="n1-standard-4"),
881-
container_spec=types.ContainerSpec(
882-
image_uri=container_uri,
883-
args=args,
884-
),
885-
)
886-
]
887-
888-
service_account = _prompt_optimizer_utils._get_service_account(config)
889-
890-
job_spec = types.CustomJobSpec(
891-
worker_pool_specs=worker_pool_specs,
892-
base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket),
893-
service_account=service_account,
894-
)
895-
896-
custom_job = types.CustomJob(
897-
display_name=display_name,
898-
job_spec=job_spec,
758+
warnings.warn(
759+
"The prompt_optimizer.optimize method is deprecated. Please use"
760+
" prompts.launch_optimization_job instead.",
761+
DeprecationWarning,
762+
stacklevel=2,
899763
)
764+
prompts_module = prompts.AsyncPrompts(api_client_=self._api_client)
900765

901-
job = await self._create_custom_job_resource(
902-
custom_job=custom_job,
766+
return await prompts_module.launch_optimization_job(
767+
method=method, config=config
903768
)
904769

905-
# Get the job id for the dashboard url and display to the user.
906-
job_resource_name = job.name
907-
if not job_resource_name:
908-
raise ValueError(f"Error creating job: {job}")
909-
job_id = job_resource_name.split("/")[-1]
910-
logger.info("Job created: %s", job.name)
911-
912-
# Construct the dashboard URL to show to the user.
913-
dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}"
914-
logger.info("View the job status at: %s", dashboard_url)
915-
916-
return job
917-
918770
async def _custom_optimize_prompt(
919771
self,
920772
*,
@@ -1021,27 +873,12 @@ async def optimize_prompt(
1021873
Returns:
1022874
The parsed response from the API request.
1023875
"""
1024-
if isinstance(config, dict):
1025-
config = types.OptimizeConfig(**config)
1026-
1027-
optimization_target: Optional[types.OptimizeTarget] = None
1028-
if config is not None:
1029-
optimization_target = config.optimization_target
1030-
1031-
final_prompt = prompt
1032-
if (
1033-
optimization_target
1034-
== types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
1035-
or optimization_target
1036-
== types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
1037-
):
1038-
final_prompt = _prompt_optimizer_utils._get_few_shot_prompt(prompt, config)
1039-
1040-
# TODO: b/435653980 - replace the custom method with a generated method.
1041-
config_for_api = config.model_copy() if config else None
1042-
return await self._custom_optimize_prompt(
1043-
content=genai_types.Content(
1044-
parts=[genai_types.Part(text=final_prompt)], role="user"
1045-
),
1046-
config=config_for_api,
876+
warnings.warn(
877+
"The prompt_optimizer.optimize_prompt method is deprecated. Please use"
878+
" prompts.optimize instead.",
879+
DeprecationWarning,
880+
stacklevel=2,
1047881
)
882+
prompts_module = prompts.AsyncPrompts(api_client_=self._api_client)
883+
884+
return await prompts_module.optimize(prompt=prompt, config=config)

0 commit comments

Comments
 (0)