google-cloud-agentplatform 1.165.1.dev0__py3-none-any.whl

agentplatform/_genai/prompts.py

@@ -0,0 +1,4515 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# Code generated by the Google Gen AI SDK generator DO NOT EDIT.
+
+import asyncio
+import datetime
+import json
+import logging
+import time
+from typing import Any, AsyncIterator, Iterator, Optional, Union
+from urllib.parse import urlencode
+
+from google.genai import _api_module
+from google.genai import _common
+from google.genai import operations
+from google.genai import types as genai_types
+from google.genai._common import get_value_by_path as getv
+from google.genai._common import set_value_by_path as setv
+from google.genai.pagers import AsyncPager, Pager
+from pydantic import ValidationError
+
+from . import _prompt_management_utils
+from . import _prompt_optimizer_utils
+from . import types
+
+logger = logging.getLogger("agentplatform_genai.prompts")
+
+
+def _CreateDatasetParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["name"]) is not None:
+        setv(to_object, ["name"], getv(from_object, ["name"]))
+
+    if getv(from_object, ["display_name"]) is not None:
+        setv(to_object, ["displayName"], getv(from_object, ["display_name"]))
+
+    if getv(from_object, ["metadata_schema_uri"]) is not None:
+        setv(
+            to_object, ["metadataSchemaUri"], getv(from_object, ["metadata_schema_uri"])
+        )
+
+    if getv(from_object, ["metadata"]) is not None:
+        setv(to_object, ["metadata"], getv(from_object, ["metadata"]))
+
+    if getv(from_object, ["description"]) is not None:
+        setv(to_object, ["description"], getv(from_object, ["description"]))
+
+    if getv(from_object, ["encryption_spec"]) is not None:
+        setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
+
+    if getv(from_object, ["model_reference"]) is not None:
+        setv(to_object, ["modelReference"], getv(from_object, ["model_reference"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _CreateDatasetVersionParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["dataset_name"]) is not None:
+        setv(to_object, ["_url", "name"], getv(from_object, ["dataset_name"]))
+
+    if getv(from_object, ["metadata"]) is not None:
+        setv(to_object, ["metadata"], getv(from_object, ["metadata"]))
+
+    if getv(from_object, ["model_reference"]) is not None:
+        setv(to_object, ["modelReference"], getv(from_object, ["model_reference"]))
+
+    if getv(from_object, ["parent"]) is not None:
+        setv(to_object, ["parent"], getv(from_object, ["parent"]))
+
+    if getv(from_object, ["display_name"]) is not None:
+        setv(to_object, ["displayName"], getv(from_object, ["display_name"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _CustomJobParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["custom_job"]) is not None:
+        setv(
+            parent_object,
+            ["customJob"],
+            _CustomJob_to_vertex(getv(from_object, ["custom_job"]), to_object),
+        )
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _CustomJob_from_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(parent_object, ["displayName"]) is not None:
+        setv(to_object, ["display_name"], getv(parent_object, ["displayName"]))
+
+    if getv(parent_object, ["jobSpec"]) is not None:
+        setv(to_object, ["job_spec"], getv(parent_object, ["jobSpec"]))
+
+    if getv(parent_object, ["encryptionSpec"]) is not None:
+        setv(to_object, ["encryption_spec"], getv(parent_object, ["encryptionSpec"]))
+
+    if getv(from_object, ["state"]) is not None:
+        setv(to_object, ["state"], getv(from_object, ["state"]))
+
+    if getv(parent_object, ["error"]) is not None:
+        setv(to_object, ["error"], getv(parent_object, ["error"]))
+
+    if getv(from_object, ["createTime"]) is not None:
+        setv(to_object, ["create_time"], getv(from_object, ["createTime"]))
+
+    if getv(from_object, ["endTime"]) is not None:
+        setv(to_object, ["end_time"], getv(from_object, ["endTime"]))
+
+    if getv(from_object, ["labels"]) is not None:
+        setv(to_object, ["labels"], getv(from_object, ["labels"]))
+
+    if getv(from_object, ["name"]) is not None:
+        setv(to_object, ["name"], getv(from_object, ["name"]))
+
+    if getv(from_object, ["satisfiesPzi"]) is not None:
+        setv(to_object, ["satisfies_pzi"], getv(from_object, ["satisfiesPzi"]))
+
+    if getv(from_object, ["satisfiesPzs"]) is not None:
+        setv(to_object, ["satisfies_pzs"], getv(from_object, ["satisfiesPzs"]))
+
+    if getv(from_object, ["startTime"]) is not None:
+        setv(to_object, ["start_time"], getv(from_object, ["startTime"]))
+
+    if getv(from_object, ["updateTime"]) is not None:
+        setv(to_object, ["update_time"], getv(from_object, ["updateTime"]))
+
+    if getv(from_object, ["webAccessUris"]) is not None:
+        setv(to_object, ["web_access_uris"], getv(from_object, ["webAccessUris"]))
+
+    return to_object
+
+
+def _CustomJob_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["display_name"]) is not None:
+        setv(parent_object, ["displayName"], getv(from_object, ["display_name"]))
+
+    if getv(from_object, ["job_spec"]) is not None:
+        setv(parent_object, ["jobSpec"], getv(from_object, ["job_spec"]))
+
+    if getv(from_object, ["encryption_spec"]) is not None:
+        setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
+
+    if getv(from_object, ["state"]) is not None:
+        setv(to_object, ["state"], getv(from_object, ["state"]))
+
+    if getv(from_object, ["error"]) is not None:
+        setv(parent_object, ["error"], getv(from_object, ["error"]))
+
+    if getv(from_object, ["create_time"]) is not None:
+        setv(to_object, ["createTime"], getv(from_object, ["create_time"]))
+
+    if getv(from_object, ["end_time"]) is not None:
+        setv(to_object, ["endTime"], getv(from_object, ["end_time"]))
+
+    if getv(from_object, ["labels"]) is not None:
+        setv(to_object, ["labels"], getv(from_object, ["labels"]))
+
+    if getv(from_object, ["name"]) is not None:
+        setv(to_object, ["name"], getv(from_object, ["name"]))
+
+    if getv(from_object, ["satisfies_pzi"]) is not None:
+        setv(to_object, ["satisfiesPzi"], getv(from_object, ["satisfies_pzi"]))
+
+    if getv(from_object, ["satisfies_pzs"]) is not None:
+        setv(to_object, ["satisfiesPzs"], getv(from_object, ["satisfies_pzs"]))
+
+    if getv(from_object, ["start_time"]) is not None:
+        setv(to_object, ["startTime"], getv(from_object, ["start_time"]))
+
+    if getv(from_object, ["update_time"]) is not None:
+        setv(to_object, ["updateTime"], getv(from_object, ["update_time"]))
+
+    if getv(from_object, ["web_access_uris"]) is not None:
+        setv(to_object, ["webAccessUris"], getv(from_object, ["web_access_uris"]))
+
+    return to_object
+
+
+def _DeleteDatasetRequestParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["prompt_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["prompt_id"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _DeletePromptVersionRequestParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["prompt_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["prompt_id"]))
+
+    if getv(from_object, ["version_id"]) is not None:
+        setv(to_object, ["_url", "version_id"], getv(from_object, ["version_id"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _GetCustomJobParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["name"]) is not None:
+        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _GetDatasetOperationParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["dataset_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["dataset_id"]))
+
+    if getv(from_object, ["operation_id"]) is not None:
+        setv(to_object, ["_url", "operation_id"], getv(from_object, ["operation_id"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _GetDatasetParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["name"]) is not None:
+        setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _GetDatasetVersionParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["dataset_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["dataset_id"]))
+
+    if getv(from_object, ["dataset_version_id"]) is not None:
+        setv(
+            to_object,
+            ["_url", "dataset_version_id"],
+            getv(from_object, ["dataset_version_id"]),
+        )
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _ListDatasetVersionsRequestParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["read_mask"]) is not None:
+        setv(to_object, ["_url", "read_mask"], getv(from_object, ["read_mask"]))
+
+    if getv(from_object, ["dataset_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["dataset_id"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(
+            to_object,
+            ["config"],
+            _ListPromptsConfig_to_vertex(getv(from_object, ["config"]), to_object),
+        )
+
+    return to_object
+
+
+def _ListDatasetsRequestParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["config"]) is not None:
+        setv(
+            to_object,
+            ["config"],
+            _ListPromptsConfig_to_vertex(getv(from_object, ["config"]), to_object),
+        )
+
+    return to_object
+
+
+def _ListPromptsConfig_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+
+    if getv(from_object, ["page_size"]) is not None:
+        setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"]))
+
+    if getv(from_object, ["page_token"]) is not None:
+        setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"]))
+
+    if getv(from_object, ["filter"]) is not None:
+        setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"]))
+
+    return to_object
+
+
+def _OptimizeConfig_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+
+    if getv(from_object, ["optimization_target"]) is not None:
+        setv(
+            parent_object,
+            ["optimizationTarget"],
+            getv(from_object, ["optimization_target"]),
+        )
+
+    return to_object
+
+
+def _OptimizeRequestParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["content"]) is not None:
+        setv(to_object, ["content"], getv(from_object, ["content"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(
+            to_object,
+            ["config"],
+            _OptimizeConfig_to_vertex(getv(from_object, ["config"]), to_object),
+        )
+
+    return to_object
+
+
+def _RestoreVersionRequestParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["dataset_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["dataset_id"]))
+
+    if getv(from_object, ["version_id"]) is not None:
+        setv(to_object, ["_url", "version_id"], getv(from_object, ["version_id"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+def _UpdateDatasetParameters_to_vertex(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+    to_object: dict[str, Any] = {}
+    if getv(from_object, ["name"]) is not None:
+        setv(to_object, ["name"], getv(from_object, ["name"]))
+
+    if getv(from_object, ["dataset_id"]) is not None:
+        setv(to_object, ["_url", "dataset_id"], getv(from_object, ["dataset_id"]))
+
+    if getv(from_object, ["display_name"]) is not None:
+        setv(to_object, ["displayName"], getv(from_object, ["display_name"]))
+
+    if getv(from_object, ["metadata"]) is not None:
+        setv(to_object, ["metadata"], getv(from_object, ["metadata"]))
+
+    if getv(from_object, ["description"]) is not None:
+        setv(to_object, ["description"], getv(from_object, ["description"]))
+
+    if getv(from_object, ["encryption_spec"]) is not None:
+        setv(to_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"]))
+
+    if getv(from_object, ["model_reference"]) is not None:
+        setv(to_object, ["modelReference"], getv(from_object, ["model_reference"]))
+
+    if getv(from_object, ["config"]) is not None:
+        setv(to_object, ["config"], getv(from_object, ["config"]))
+
+    return to_object
+
+
+class Prompts(_api_module.BaseModule):
+
+    def _create_dataset_resource(
+        self,
+        *,
+        name: Optional[str] = None,
+        display_name: Optional[str] = None,
+        metadata_schema_uri: Optional[str] = None,
+        metadata: Optional[types.SchemaTextPromptDatasetMetadataOrDict] = None,
+        description: Optional[str] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
+        model_reference: Optional[str] = None,
+        config: Optional[types.CreateDatasetConfigOrDict] = None,
+    ) -> types.DatasetOperation:
+        """
+        Creates a dataset resource to store prompts.
+        """
+
+        parameter_model = types._CreateDatasetParameters(
+            name=name,
+            display_name=display_name,
+            metadata_schema_uri=metadata_schema_uri,
+            metadata=metadata,
+            description=description,
+            encryption_spec=encryption_spec,
+            model_reference=model_reference,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _CreateDatasetParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets".format_map(request_url_dict)
+            else:
+                path = "datasets"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("post", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _create_dataset_version_resource(
+        self,
+        *,
+        dataset_name: Optional[str] = None,
+        metadata: Optional[types.SchemaTextPromptDatasetMetadataOrDict] = None,
+        model_reference: Optional[str] = None,
+        parent: Optional[str] = None,
+        display_name: Optional[str] = None,
+        config: Optional[types.CreateDatasetVersionConfigOrDict] = None,
+    ) -> types.DatasetOperation:
+        """
+        Creates a dataset version resource to store prompts.
+        """
+
+        parameter_model = types._CreateDatasetVersionParameters(
+            dataset_name=dataset_name,
+            metadata=metadata,
+            model_reference=model_reference,
+            parent=parent,
+            display_name=display_name,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _CreateDatasetVersionParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{name}/datasetVersions".format_map(request_url_dict)
+            else:
+                path = "datasets/{name}/datasetVersions"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("post", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _get_dataset_resource(
+        self,
+        *,
+        name: Optional[str] = None,
+        config: Optional[types.VertexBaseConfigOrDict] = None,
+    ) -> types.Dataset:
+        """
+        Gets a dataset resource to store prompts.
+        """
+
+        parameter_model = types._GetDatasetParameters(
+            name=name,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetDatasetParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{name}".format_map(request_url_dict)
+            else:
+                path = "datasets/{name}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.Dataset._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _get_dataset_version_resource(
+        self,
+        *,
+        dataset_id: Optional[str] = None,
+        dataset_version_id: Optional[str] = None,
+        config: Optional[types.VertexBaseConfigOrDict] = None,
+    ) -> types.DatasetVersion:
+        """
+        Gets a dataset version resource to store prompts.
+        """
+
+        parameter_model = types._GetDatasetVersionParameters(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetDatasetVersionParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions/{dataset_version_id}".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions/{dataset_version_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetVersion._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _get_dataset_operation(
+        self,
+        *,
+        dataset_id: Optional[str] = None,
+        operation_id: Optional[str] = None,
+        config: Optional[types.GetDatasetOperationConfigOrDict] = None,
+    ) -> types.DatasetOperation:
+        """
+        Gets the operation from creating a dataset.
+        """
+
+        parameter_model = types._GetDatasetOperationParameters(
+            dataset_id=dataset_id,
+            operation_id=operation_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetDatasetOperationParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/operations/{operation_id}".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/operations/{operation_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _list_prompts(
+        self, *, config: Optional[types.ListPromptsConfigOrDict] = None
+    ) -> types.ListDatasetsResponse:
+        """
+        Lists Agent Engines.
+        """
+
+        parameter_model = types._ListDatasetsRequestParameters(
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _ListDatasetsRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets".format_map(request_url_dict)
+            else:
+                path = "datasets"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.ListDatasetsResponse._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _list_versions(
+        self,
+        *,
+        read_mask: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> types.ListDatasetVersionsResponse:
+        """
+        Lists Agent Engines.
+        """
+
+        parameter_model = types._ListDatasetVersionsRequestParameters(
+            read_mask=read_mask,
+            dataset_id=dataset_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _ListDatasetVersionsRequestParameters_to_vertex(
+                parameter_model
+            )
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.ListDatasetVersionsResponse._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _delete_dataset(
+        self, *, prompt_id: str, config: Optional[types.DeletePromptConfigOrDict] = None
+    ) -> types.DeletePromptOperation:
+        parameter_model = types._DeleteDatasetRequestParameters(
+            prompt_id=prompt_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _DeleteDatasetRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}".format_map(request_url_dict)
+            else:
+                path = "datasets/{dataset_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("delete", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DeletePromptOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _delete_dataset_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.DeletePromptConfigOrDict] = None,
+    ) -> types.DeletePromptVersionOperation:
+        parameter_model = types._DeletePromptVersionRequestParameters(
+            prompt_id=prompt_id,
+            version_id=version_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _DeletePromptVersionRequestParameters_to_vertex(
+                parameter_model
+            )
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("delete", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DeletePromptVersionOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _restore_version(
+        self,
+        *,
+        dataset_id: str,
+        version_id: str,
+        config: Optional[types.RestoreVersionConfigOrDict] = None,
+    ) -> types.RestoreVersionOperation:
+        """
+        Restores the provided prompt version to the latest version.
+        """
+
+        parameter_model = types._RestoreVersionRequestParameters(
+            dataset_id=dataset_id,
+            version_id=version_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _RestoreVersionRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}:restore".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}:restore"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.RestoreVersionOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _update_dataset_resource(
+        self,
+        *,
+        name: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        display_name: Optional[str] = None,
+        metadata: Optional[types.SchemaTextPromptDatasetMetadataOrDict] = None,
+        description: Optional[str] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
+        model_reference: Optional[str] = None,
+        config: Optional[types.UpdatePromptConfigOrDict] = None,
+    ) -> types.Dataset:
+        """
+        Creates a dataset resource to store prompts.
+        """
+
+        parameter_model = types._UpdateDatasetParameters(
+            name=name,
+            dataset_id=dataset_id,
+            display_name=display_name,
+            metadata=metadata,
+            description=description,
+            encryption_spec=encryption_spec,
+            model_reference=model_reference,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _UpdateDatasetParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}".format_map(request_url_dict)
+            else:
+                path = "datasets/{dataset_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("patch", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.Dataset._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _create_custom_job_resource(
+        self,
+        *,
+        custom_job: types.CustomJobOrDict,
+        config: Optional[types.VertexBaseConfigOrDict] = None,
+    ) -> types.CustomJob:
+        """
+        Creates a custom job.
+        """
+
+        parameter_model = types._CustomJobParameters(
+            custom_job=custom_job,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _CustomJobParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "customJobs".format_map(request_url_dict)
+            else:
+                path = "customJobs"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("post", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        if self._api_client.vertexai:
+            response_dict = _CustomJob_from_vertex(response_dict)
+
+        return_value = types.CustomJob._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _get_custom_job(
+        self, *, name: str, config: Optional[types.VertexBaseConfigOrDict] = None
+    ) -> types.CustomJob:
+        """
+        Gets a custom job.
+        """
+
+        parameter_model = types._GetCustomJobParameters(
+            name=name,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetCustomJobParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "customJobs/{name}".format_map(request_url_dict)
+            else:
+                path = "customJobs/{name}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("get", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        if self._api_client.vertexai:
+            response_dict = _CustomJob_from_vertex(response_dict)
+
+        return_value = types.CustomJob._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    def _optimize(
+        self,
+        *,
+        content: Optional[genai_types.ContentOrDict] = None,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> types.OptimizeResponseEndpoint:
+        """
+        Optimize a single prompt.
+        """
+
+        parameter_model = types._OptimizeRequestParameters(
+            content=content,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
+            else:
+                path = "tuningJobs:optimizePrompt"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("post", path, request_dict, http_options)
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.OptimizeResponseEndpoint._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    _DEFAULT_TIMEOUT = 90
+    _DEFAULT_MAX_WAIT_TIME = 60
+
+    def create(
+        self,
+        *,
+        prompt: types.PromptOrDict,
+        config: Optional[types.CreatePromptConfigOrDict] = None,
+    ) -> types.Prompt:
+        """Creates a new prompt in a Vertex Dataset resource.
+
+        This method waits for prompt creation to be complete before returning.
+
+        Note: This method does not create a versioned resource for your prompt.
+        Call create_version() to create a new prompt resource with a version.
+
+        Args:
+          prompt: The prompt to create.
+          config: Optional configuration for creating the prompt.
+
+        Returns:
+          A types.Prompt object representing the prompt with its associated
+          Dataset resources.
+        """
+        if isinstance(prompt, dict):
+            prompt = types.Prompt(**prompt)
+        if isinstance(config, dict):
+            config = types.CreatePromptConfig(**config)
+        elif not config:
+            config = types.CreatePromptConfig()
+
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
+
+        if prompt.prompt_data is None:
+            raise ValueError("Prompt data is required to create a prompt.")
+
+        prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
+            prompt,
+            variables=(
+                prompt.prompt_data.variables
+                if prompt.prompt_data and prompt.prompt_data.variables
+                else None
+            ),
+        )
+
+        # Step 1: Create the dataset resource for the prompt and wait for the operation to complete.
+        create_prompt_dataset_operation = self._create_dataset_resource(
+            display_name=(
+                config.prompt_display_name
+                if config and config.prompt_display_name
+                else f"prompt_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}",
+            metadata_schema_uri=_prompt_management_utils.PROMPT_SCHEMA_URI,
+            metadata=prompt_metadata,
+            model_reference=prompt.prompt_data.model,
+            encryption_spec=(
+                config.encryption_spec if config and config.encryption_spec else None
+            ),
+        )
+        dataset_resource_name = self._wait_for_operation(
+            operation=create_prompt_dataset_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_id = dataset_resource_name.split("/")[-1]
+
+        # Step 2: Get the dataset resource
+        dataset_resource = self._get_dataset_resource(
+            name=dataset_id,
+        )
+        prompt._dataset = dataset_resource
+        return prompt
+
+    def create_version(
+        self,
+        *,
+        prompt: types.PromptOrDict,
+        config: Optional[types.CreatePromptVersionConfigOrDict] = None,
+    ) -> types.Prompt:
+        """Creates a prompt resource and an initial prompt version.
+
+        When creating new prompt and prompt version resources, this waits for
+        the create operation to complete before returning.
+
+        Note: This method is recommended instead of create() since it creates a
+        versioned resource for your prompt.
+
+        Args:
+          prompt: The prompt to create.
+          config: Optional configuration for creating the prompt and prompt version.
+
+        Returns:
+          A types.Prompt object representing the prompt with its associated
+          Dataset and Dataset Version resources.
+        """
+
+        if isinstance(prompt, dict):
+            prompt = types.Prompt(**prompt)
+        if isinstance(config, dict):
+            config = types.CreatePromptVersionConfig(**config)
+        elif not config:
+            config = types.CreatePromptVersionConfig()
+
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
+
+        if prompt.prompt_data is None:
+            raise ValueError("Prompt data is required to create a prompt.")
+
+        prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
+            prompt,
+            variables=(
+                prompt.prompt_data.variables
+                if prompt.prompt_data and prompt.prompt_data.variables
+                else None
+            ),
+        )
+
+        # Step 1: Create the dataset resource for the prompt and wait for the operation to complete.
+        create_prompt_dataset_operation = self._create_dataset_resource(
+            display_name=(
+                config.prompt_display_name
+                if config and config.prompt_display_name
+                else f"prompt_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}",
+            metadata_schema_uri=_prompt_management_utils.PROMPT_SCHEMA_URI,
+            metadata=prompt_metadata,
+            model_reference=prompt.prompt_data.model,
+            encryption_spec=(
+                config.encryption_spec if config and config.encryption_spec else None
+            ),
+        )
+        dataset_resource_name = self._wait_for_operation(
+            operation=create_prompt_dataset_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_id = dataset_resource_name.split("/")[-1]
+
+        # Step 2: Get the dataset resource
+        dataset_resource = self._get_dataset_resource(
+            name=dataset_id,
+        )
+        prompt._dataset = dataset_resource
+
+        if prompt._dataset.name is None:
+            raise ValueError("Invalid dataset resource.")
+
+        # Step 3: Create the dataset version
+        create_dataset_version_operation = self._create_dataset_version_resource(
+            dataset_name=prompt._dataset.name.split("/")[-1],
+            display_name=(
+                config.version_display_name
+                if config and config.version_display_name is not None
+                else f"prompt_version_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+        )
+        dataset_version_resource_name = self._wait_for_operation(
+            operation=create_dataset_version_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+
+        # Step 4: Get the dataset version resource and return it with the prompt
+        dataset_version_resource = self._get_dataset_version_resource(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_resource_name.split("/")[-1],
+        )
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource
+        )
+        prompt._dataset = dataset_resource
+        prompt._dataset_version = dataset_version_resource
+        return prompt
+
+    def _wait_for_operation(
+        self,
+        operation: types.DatasetOperation,
+        timeout: int,
+        max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
+    ) -> str:
+        """Waits for a dataset operation to complete.
+
+        Args:
+          operation: The dataset operation to wait for.
+          timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
+
+        Returns:
+          The name of the Dataset resource from the operation result.
+
+        Raises:
+          TimeoutError: If the operation does not complete within the timeout.
+          ValueError: If the operation fails.
+        """
+        done = False
+        prompt_dataset_operation: Optional[types.DatasetOperation] = None
+
+        response_operation_name = operation.name
+        if response_operation_name is None:
+            raise ValueError("Invalid operation name.")
+
+        dataset_id = response_operation_name.split("/datasets/")[1].split("/")[0]
+        operation_id = response_operation_name.split("/")[-1]
+
+        start_time = time.time()
+        sleep_duration = 5
+        wait_multiplier = 2
+        previous_time = time.time()
+
+        while not done:
+            if (time.time() - start_time) > timeout:
+                raise TimeoutError(
+                    "Create prompt operation did not complete within the"
+                    f" specified timeout of {timeout} seconds."
+                )
+            current_time = time.time()
+            if current_time - previous_time >= sleep_duration:
+                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
+                previous_time = current_time
+            time.sleep(sleep_duration)
+            prompt_dataset_operation = self._get_dataset_operation(
+                dataset_id=dataset_id,
+                operation_id=operation_id,
+            )
+            done = (
+                (prompt_dataset_operation.done or False)
+                if hasattr(prompt_dataset_operation, "done")
+                else False
+            )
+        if (
+            not prompt_dataset_operation
+            or prompt_dataset_operation.response is None
+            or prompt_dataset_operation.response.get("name") is None
+        ):
+            raise ValueError("Error creating prompt version resource.")
+        if (
+            hasattr(prompt_dataset_operation, "error")
+            and prompt_dataset_operation.error is not None
+        ):
+            raise ValueError(
+                f"Error creating prompt version resource: {prompt_dataset_operation.error}"
+            )
+        return prompt_dataset_operation.response.get("name")  # type: ignore[return-value]
+
+    def get(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.GetPromptConfig] = None,
+    ) -> types.Prompt:
+        """Gets a prompt resource from a Vertex Dataset.
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          config: Optional configuration for getting the prompt.
+
+        Returns:
+            A types.Prompt object representing the prompt with its associated Dataset resources.
+        """
+
+        prompt_dataset_resource = self._get_dataset_resource(name=prompt_id)
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            prompt_dataset_resource,
+        )
+        prompt._dataset = prompt_dataset_resource
+
+        return prompt
+
+    def get_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.GetPromptConfig] = None,
+    ) -> types.Prompt:
+        """Gets a prompt resource from a Vertex Dataset.
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          version_id: The id of the Vertex Dataset Version resource containing the prompt version. For example, if the prompt version resource name is "projects/123/locations/us-central1/datasets/456/datasetVersions/1", then the version_id is "1".
+          config: Optional configuration for getting the prompt.
+
+        Returns:
+            A types.Prompt object representing the prompt with its associated Dataset and Dataset Version resources.
+        """
+
+        prompt_dataset_resource = self._get_dataset_resource(name=prompt_id)
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            prompt_dataset_resource,
+        )
+        prompt._dataset = prompt_dataset_resource
+
+        prompt_version_resource = self._get_dataset_version_resource(
+            dataset_id=prompt_id,
+            dataset_version_id=version_id,
+        )
+        prompt._dataset_version = prompt_version_resource
+
+        return prompt
+
+    def _list_prompts_pager(
+        self,
+        *,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> Pager[types.Dataset]:
+        return Pager(
+            "datasets",
+            self._list_prompts,
+            self._list_prompts(config=config),
+            config,
+        )
+
+    def _list_versions_pager(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> Pager[types.DatasetVersion]:
+        return Pager(
+            "dataset_versions",
+            self._list_versions,
+            self._list_versions(config=config, dataset_id=prompt_id),
+            config,
+        )
+
+    def list(
+        self,
+        *,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> Iterator[types.PromptRef]:
+        """Lists prompt resources in a project.
+
+        This method retrieves all the prompts from the project provided in the
+        agentplatform.Client constructor and returns a list of prompt references containing the prompt_id and model for the prompt.
+
+        To get the full types.Prompt resource for a PromptRef after calling this method, use the get() method with the prompt_id as the prompt_id argument.
+        Example usage:
+
+        ```
+        # Using an iterator
+        prompt_refs = client.prompt_management.list_prompts()
+        for prompt_ref in prompt_refs:
+          client.prompt_management.get(prompt_id=prompt_ref.prompt_id)
+
+        # Using a list
+        prompts_list = list(client.prompt_management.list_prompts())
+        client.prompt_management.get(prompt_id=prompts_list[0].prompt_id)
+        ```
+
+        Args:
+          config: Optional configuration for listing prompts.
+
+        Returns:
+            An iterable of types.PromptRef objects.
+        """
+        if isinstance(config, dict):
+            config = types.ListPromptsConfig(**config)
+        elif not config:
+            config = types.ListPromptsConfig()
+        for dataset in self._list_prompts_pager(config=config):
+            if not dataset.name:
+                continue
+            prompt_ref = types.PromptRef(
+                model=dataset.model_reference, prompt_id=dataset.name.split("/")[-1]
+            )
+            yield prompt_ref
+
+    def list_versions(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> Iterator[types.PromptVersionRef]:
+        """Lists prompt version resources for a provided prompt_id.
+
+        This method retrieves all the prompt versions for a provided prompt_id.
+
+        To get the full types.Prompt resource for a PromptVersionRef after calling this method, use the get() method with the returned prompt_id and version_id.
+        Example usage:
+
+        ```
+        # Using an iterator
+        prompt_version_refs = client.prompt_management.list_versions(prompt_id="123")
+        for version_ref in prompt_version_refs:
+          client.prompt_management.get(prompt_id=version_ref.prompt_id, version_id=version_ref.version_id)
+
+        # Using a list
+        prompt_versions_list = list(client.prompt_management.list_versions(prompt_id="123"))
+        client.prompt_management.get(prompt_id=prompt_versions_list[0].prompt_id, version_id=prompt_versions_list[0].version_id)
+        ```
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          config: Optional configuration for listing prompts.
+
+        Returns:
+            An iterable of types.PromptVersionRef objects representing the prompt version resources for the provided prompt_id.
+
+        """
+        if isinstance(config, dict):
+            config = types.ListPromptsConfig(**config)
+        elif not config:
+            config = types.ListPromptsConfig()
+        for dataset_version in self._list_versions_pager(
+            config=config, prompt_id=prompt_id
+        ):
+            if (
+                not dataset_version
+                or not dataset_version.model_reference
+                or not dataset_version.name
+            ):
+                continue
+            prompt_version_ref = types.PromptVersionRef(
+                model=dataset_version.model_reference,
+                version_id=dataset_version.name.split("/")[-1],
+                prompt_id=prompt_id,
+            )
+            yield prompt_version_ref
+
+    def _wait_for_project_operation(
+        self,
+        operation: genai_types.ProjectOperation,
+        timeout: int,
+        max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
+    ) -> None:
+        """Waits for a dataset deletion operation to complete.
+
+        Delete operations are project level operations and are separate from dataset resource operations, for example: projects/123/locations/us-central1/operations/789.
+
+        Args:
+          operation: The project operation to wait for.
+          timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
+        Raises:
+          TimeoutError: If the operation does not complete within the timeout.
+          ValueError: If the operation fails.
+        """
+        done = False
+
+        start_time = time.time()
+        sleep_duration = 5
+        wait_multiplier = 2
+        previous_time = time.time()
+        while not done:
+            if (time.time() - start_time) > timeout:
+                raise TimeoutError(
+                    f"Delete operation did not complete within the"
+                    f" specified timeout of {timeout} seconds."
+                )
+            current_time = time.time()
+            if current_time - previous_time >= sleep_duration:
+                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
+                previous_time = current_time
+            time.sleep(sleep_duration)
+            operations_module = operations.Operations(api_client_=self._api_client)
+
+            if operation.name is None:
+                raise ValueError("Invalid operation name.")
+            operation = operations_module._get(
+                operation_id=operation.name.split("/")[-1],
+            )
+            done = (operation.done or False) if hasattr(operation, "done") else False
+        if hasattr(operation, "error") and operation.error is not None:
+            raise ValueError(f"Error in delete operation: {operation.error}")
+
+    def delete(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.DeletePromptConfig] = None,
+    ) -> None:
+        """Deletes a prompt resource.
+
+        Args:
+          prompt_id: The id of the prompt resource to delete.
+
+        Raises:
+          TimeoutError: If the delete operation does not complete within the timeout.
+          ValueError: If the delete operation fails.
+        """
+
+        delete_prompt_operation = self._delete_dataset(
+            prompt_id=prompt_id,
+            config=config,
+        )
+        self._wait_for_project_operation(
+            operation=delete_prompt_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        logger.info(f"Deleted prompt with id: {prompt_id}")
+
+    def delete_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.DeletePromptConfig] = None,
+    ) -> None:
+        """Deletes a prompt version resource.
+
+        Args:
+          prompt_id: The id of the prompt resource to delete.
+          version_id: The id of the prompt version resource to delete.
+
+        Raises:
+          TimeoutError: If the delete operation does not complete within the timeout.
+          ValueError: If the delete operation fails.
+        """
+        delete_version_operation = self._delete_dataset_version(
+            prompt_id=prompt_id,
+            version_id=version_id,
+            config=config,
+        )
+
+        self._wait_for_project_operation(
+            operation=delete_version_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        logger.info(
+            f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
+        )
+
+    def restore_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.RestoreVersionConfig] = None,
+    ) -> types.Prompt:
+        """Restores the provided prompt version to the latest version.
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          version_id: The id of the Vertex Dataset Version resource to restore. For example, if the version resource name is "projects/123/locations/us-central1/datasets/456/datasetVersions/789", then the version_id is "789".
+          config: Optional configuration for restoring the prompt version.
+
+        Returns:
+            A types.Prompt object representing the prompt with the updated Dataset Version resource.
+        """
+
+        restore_prompt_operation = self._restore_version(
+            dataset_id=prompt_id,
+            version_id=version_id,
+            config=config,
+        )
+        self._wait_for_project_operation(
+            operation=restore_prompt_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_version_resource = self._get_dataset_version_resource(
+            dataset_id=prompt_id,
+            dataset_version_id=version_id,
+        )
+        updated_prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource,
+        )
+        updated_prompt._dataset_version = dataset_version_resource
+        return updated_prompt
+
+    def _wait_for_completion(self, job_name: str) -> types.CustomJob:
+
+        JOB_COMPLETE_STATES = [
+            genai_types.JobState.JOB_STATE_SUCCEEDED,
+            genai_types.JobState.JOB_STATE_FAILED,
+            genai_types.JobState.JOB_STATE_CANCELLED,
+            genai_types.JobState.JOB_STATE_PAUSED,
+        ]
+        JOB_ERROR_STATES = [
+            genai_types.JobState.JOB_STATE_FAILED,
+            genai_types.JobState.JOB_STATE_CANCELLED,
+        ]
+
+        log_wait = 5
+        wait_multiplier = 2
+        max_wait_time = self._DEFAULT_MAX_WAIT_TIME
+        previous_time = time.time()
+
+        job = self._get_custom_job(name=job_name)
+
+        while job.state not in JOB_COMPLETE_STATES:
+            current_time = time.time()
+            if current_time - previous_time >= log_wait:
+                logger.info(f"Waiting for job to complete. Current state: {job.state}")
+                log_wait = min(log_wait * wait_multiplier, max_wait_time)
+                previous_time = current_time
+            time.sleep(log_wait)
+            job = self._get_custom_job(name=job_name)
+
+        logger.info(f"Job state: {job.state}")
+
+        if job.state in JOB_ERROR_STATES:
+            raise RuntimeError(f"Job failed with state: {job.state}")
+        else:
+            logger.info(f"Job completed with state: {job.state}")
+        return job
+
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI prompts.launch_optimization_job method is "
+        "experimental, and may change in future versions."
+    )
+    def launch_optimization_job(
+        self,
+        method: types.PromptOptimizerMethod,
+        config: types.PromptOptimizerConfigOrDict,
+    ) -> types.CustomJob:
+        """Call PO-Data optimizer.
+
+        Args:
+          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 isinstance(config, dict):
+            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/vaipo:preview_android_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"{method.value.lower()}-optimizer-{timestamp}"
+
+        wait_for_completion = config.wait_for_completion
+        bucket = "/".join(config.config_path.split("/")[:-1])
+
+        region = self._api_client.location
+        project = self._api_client.project
+        container_args = {
+            "config": config.config_path,
+        }
+        args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
+        worker_pool_specs = [
+            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)
+
+        job_spec = types.CustomJobSpec(
+            worker_pool_specs=worker_pool_specs,
+            base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket),
+            service_account=service_account,
+        )
+
+        custom_job = types.CustomJob(
+            display_name=display_name,
+            job_spec=job_spec,
+        )
+
+        job = self._create_custom_job_resource(
+            custom_job=custom_job,
+        )
+
+        # Get the job resource name
+        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)
+
+        # Construct the dashboard URL
+        dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}"
+        logger.info("View the job status at: %s", dashboard_url)
+
+        if wait_for_completion:
+            job = self._wait_for_completion(job_id)
+        return job
+
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI prompts.optimize method is "
+        "experimental, and may change in future versions."
+    )
+    def optimize(
+        self,
+        *,
+        prompt: str,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> types.OptimizeResponse:
+        """Makes an API request to optimize a prompt and returns the parsed response.
+
+        Example usage:
+        client = agentplatform.Client(project=PROJECT_NAME, location='us-central1')
+        prompt = "Generate system instructions for analyzing medical articles"
+        response = client.prompts.optimize(prompt=prompt)
+
+        Args:
+          prompt: Required. The prompt to optimize.
+          config: Optional. The configuration for prompt optimization. To optimize
+            prompts from Android API provide
+            types.OptimizeConfig(
+                optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO
+            )
+            For few-shot optimization, provide:
+
+            optim_target = types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
+            or
+            optim_target = types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
+            types.OptimizeConfig(
+                optimization_target=optim_target,
+                examples_dataframe=dataframe
+            )
+            OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS indicates that the few-shot
+              examples include specific scoring rubrics and their corresponding
+              evaluations.
+            OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE indicates that the few-shot
+              examples include a ground-truth target response.
+        Returns:
+          The parsed response from the API request.
+        """
+
+        if isinstance(config, dict):
+            config = types.OptimizeConfig(**config)
+
+        optimization_target: Optional[types.OptimizeTarget] = None
+        if config is not None:
+            optimization_target = config.optimization_target
+
+        final_prompt = prompt
+        if (
+            optimization_target
+            == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
+            or optimization_target
+            == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
+        ):
+            final_prompt = _prompt_optimizer_utils._get_few_shot_prompt(prompt, config)
+
+        # TODO: b/435653980 - replace the custom method with a generated method.
+        config_for_api = config.model_copy() if config else None
+        return self._custom_optimize(
+            content=genai_types.Content(
+                parts=[genai_types.Part(text=final_prompt)], role="user"
+            ),
+            config=config_for_api,
+        )
+
+    def _custom_optimize(
+        self,
+        *,
+        content: Optional[genai_types.ContentOrDict] = None,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> types.OptimizeResponse:
+        """Internal method to call the optimizePrompt endpoint.
+
+        Sends a request to the tuningJobs:optimizePrompt streaming endpoint.
+        Then gathers the response, concatenates into one string and returns
+        the parsed response.
+        """
+        if isinstance(config, dict):
+            config.pop("examples_dataframe", None)
+        elif config and hasattr(config, "examples_dataframe"):
+            del config.examples_dataframe
+
+        parameter_model = types._OptimizeRequestParameters(
+            content=content,
+            config=config,
+        )
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError("This method is only supported in the Vertex AI client.")
+        else:
+            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
+            else:
+                path = "tuningJobs:optimizePrompt"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[genai_types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = self._api_client.request("post", path, request_dict, http_options)
+
+        response_list = "" if not response.body else json.loads(response.body)
+
+        return_value = []
+
+        for response_dict in response_list:
+            response_value = types.OptimizeResponseEndpoint._from_response(
+                response=response_dict, kwargs=parameter_model.model_dump()
+            )
+            self._api_client._verify_response(response_value)
+            content = response_value.content
+            if content is not None:
+                parts = content.parts
+                if parts and parts[0].text is not None:
+                    return_value.append(parts[0].text)
+
+        output = "".join(return_value)
+        final_response = types.OptimizeResponse(raw_text_response=output)
+        try:
+            final_response.parsed_response = _prompt_optimizer_utils._parse(output)
+        except (ValueError, TypeError, ValidationError) as e:
+            logger.warning(
+                f"Failed to parse response: {e}. Returning only raw_text_response."
+            )
+        return final_response
+
+    def update(
+        self,
+        *,
+        prompt_id: str,
+        prompt: types.PromptOrDict,
+        config: Optional[types.UpdatePromptConfigOrDict] = None,
+    ) -> types.Prompt:
+        """Updates an existing prompt and creates a new version for the prompt associated with the provided prompt_id.
+
+        Args:
+          prompt_id: The ID of the prompt to create a new version for.
+          prompt: The updated prompt.
+          config: Optional configuration for updating the prompt.
+
+        Returns:
+          A types.Prompt object representing the updated prompt with its associated
+          Dataset and Dataset Version resources.
+        """
+
+        if isinstance(prompt, dict):
+            prompt = types.Prompt(**prompt)
+        if isinstance(config, dict):
+            config = types.UpdatePromptConfig(**config)
+        elif not config:
+            config = types.UpdatePromptConfig()
+
+        prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
+            prompt,
+            variables=(
+                prompt.prompt_data.variables
+                if prompt.prompt_data and prompt.prompt_data.variables
+                else None
+            ),
+        )
+
+        if not prompt.prompt_data:
+            raise ValueError("Prompt data is required to update a prompt.")
+
+        # Step 1: Update the dataset resource for the prompt and wait for the operation to complete.
+        updated_dataset_resource = self._update_dataset_resource(
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/{prompt_id}",
+            dataset_id=prompt_id,
+            display_name=(
+                config.prompt_display_name
+                if config and config.prompt_display_name
+                else None
+            ),
+            metadata=prompt_metadata,
+            model_reference=prompt.prompt_data.model,
+            encryption_spec=(
+                config.encryption_spec if config and config.encryption_spec else None
+            ),
+            config=config,
+        )
+
+        if not updated_dataset_resource.name:
+            raise ValueError("Failed to update dataset resource.")
+
+        dataset_id = updated_dataset_resource.name.split("/")[-1]
+
+        # Step 2: Create a dataset version for the prompt.
+        create_dataset_version_operation = self._create_dataset_version_resource(
+            dataset_name=dataset_id,
+            display_name=(
+                config.version_display_name
+                if config and config.version_display_name is not None
+                else f"prompt_version_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+        )
+        dataset_version_resource_name = self._wait_for_operation(
+            operation=create_dataset_version_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_version_id = dataset_version_resource_name.split("/")[-1]
+
+        # Step 3: Get the dataset version resource and return it with the prompt.
+        dataset_version_resource = self._get_dataset_version_resource(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_id,
+        )
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource
+        )
+        prompt._dataset = updated_dataset_resource
+        prompt._dataset_version = dataset_version_resource
+        return prompt
+
+
+class AsyncPrompts(_api_module.BaseModule):
+
+    async def _create_dataset_resource(
+        self,
+        *,
+        name: Optional[str] = None,
+        display_name: Optional[str] = None,
+        metadata_schema_uri: Optional[str] = None,
+        metadata: Optional[types.SchemaTextPromptDatasetMetadataOrDict] = None,
+        description: Optional[str] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
+        model_reference: Optional[str] = None,
+        config: Optional[types.CreateDatasetConfigOrDict] = None,
+    ) -> types.DatasetOperation:
+        """
+        Creates a dataset resource to store prompts.
+        """
+
+        parameter_model = types._CreateDatasetParameters(
+            name=name,
+            display_name=display_name,
+            metadata_schema_uri=metadata_schema_uri,
+            metadata=metadata,
+            description=description,
+            encryption_spec=encryption_spec,
+            model_reference=model_reference,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _CreateDatasetParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets".format_map(request_url_dict)
+            else:
+                path = "datasets"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "post", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _create_dataset_version_resource(
+        self,
+        *,
+        dataset_name: Optional[str] = None,
+        metadata: Optional[types.SchemaTextPromptDatasetMetadataOrDict] = None,
+        model_reference: Optional[str] = None,
+        parent: Optional[str] = None,
+        display_name: Optional[str] = None,
+        config: Optional[types.CreateDatasetVersionConfigOrDict] = None,
+    ) -> types.DatasetOperation:
+        """
+        Creates a dataset version resource to store prompts.
+        """
+
+        parameter_model = types._CreateDatasetVersionParameters(
+            dataset_name=dataset_name,
+            metadata=metadata,
+            model_reference=model_reference,
+            parent=parent,
+            display_name=display_name,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _CreateDatasetVersionParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{name}/datasetVersions".format_map(request_url_dict)
+            else:
+                path = "datasets/{name}/datasetVersions"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "post", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _get_dataset_resource(
+        self,
+        *,
+        name: Optional[str] = None,
+        config: Optional[types.VertexBaseConfigOrDict] = None,
+    ) -> types.Dataset:
+        """
+        Gets a dataset resource to store prompts.
+        """
+
+        parameter_model = types._GetDatasetParameters(
+            name=name,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetDatasetParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{name}".format_map(request_url_dict)
+            else:
+                path = "datasets/{name}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.Dataset._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _get_dataset_version_resource(
+        self,
+        *,
+        dataset_id: Optional[str] = None,
+        dataset_version_id: Optional[str] = None,
+        config: Optional[types.VertexBaseConfigOrDict] = None,
+    ) -> types.DatasetVersion:
+        """
+        Gets a dataset version resource to store prompts.
+        """
+
+        parameter_model = types._GetDatasetVersionParameters(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetDatasetVersionParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions/{dataset_version_id}".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions/{dataset_version_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetVersion._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _get_dataset_operation(
+        self,
+        *,
+        dataset_id: Optional[str] = None,
+        operation_id: Optional[str] = None,
+        config: Optional[types.GetDatasetOperationConfigOrDict] = None,
+    ) -> types.DatasetOperation:
+        """
+        Gets the operation from creating a dataset.
+        """
+
+        parameter_model = types._GetDatasetOperationParameters(
+            dataset_id=dataset_id,
+            operation_id=operation_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetDatasetOperationParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/operations/{operation_id}".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/operations/{operation_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DatasetOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _list_prompts(
+        self, *, config: Optional[types.ListPromptsConfigOrDict] = None
+    ) -> types.ListDatasetsResponse:
+        """
+        Lists Agent Engines.
+        """
+
+        parameter_model = types._ListDatasetsRequestParameters(
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _ListDatasetsRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets".format_map(request_url_dict)
+            else:
+                path = "datasets"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.ListDatasetsResponse._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _list_versions(
+        self,
+        *,
+        read_mask: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> types.ListDatasetVersionsResponse:
+        """
+        Lists Agent Engines.
+        """
+
+        parameter_model = types._ListDatasetVersionsRequestParameters(
+            read_mask=read_mask,
+            dataset_id=dataset_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _ListDatasetVersionsRequestParameters_to_vertex(
+                parameter_model
+            )
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.ListDatasetVersionsResponse._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _delete_dataset(
+        self, *, prompt_id: str, config: Optional[types.DeletePromptConfigOrDict] = None
+    ) -> types.DeletePromptOperation:
+        parameter_model = types._DeleteDatasetRequestParameters(
+            prompt_id=prompt_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _DeleteDatasetRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}".format_map(request_url_dict)
+            else:
+                path = "datasets/{dataset_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "delete", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DeletePromptOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _delete_dataset_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.DeletePromptConfigOrDict] = None,
+    ) -> types.DeletePromptVersionOperation:
+        parameter_model = types._DeletePromptVersionRequestParameters(
+            prompt_id=prompt_id,
+            version_id=version_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _DeletePromptVersionRequestParameters_to_vertex(
+                parameter_model
+            )
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "delete", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.DeletePromptVersionOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _restore_version(
+        self,
+        *,
+        dataset_id: str,
+        version_id: str,
+        config: Optional[types.RestoreVersionConfigOrDict] = None,
+    ) -> types.RestoreVersionOperation:
+        """
+        Restores the provided prompt version to the latest version.
+        """
+
+        parameter_model = types._RestoreVersionRequestParameters(
+            dataset_id=dataset_id,
+            version_id=version_id,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _RestoreVersionRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}:restore".format_map(
+                    request_url_dict
+                )
+            else:
+                path = "datasets/{dataset_id}/datasetVersions/{version_id}:restore"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.RestoreVersionOperation._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _update_dataset_resource(
+        self,
+        *,
+        name: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        display_name: Optional[str] = None,
+        metadata: Optional[types.SchemaTextPromptDatasetMetadataOrDict] = None,
+        description: Optional[str] = None,
+        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
+        model_reference: Optional[str] = None,
+        config: Optional[types.UpdatePromptConfigOrDict] = None,
+    ) -> types.Dataset:
+        """
+        Creates a dataset resource to store prompts.
+        """
+
+        parameter_model = types._UpdateDatasetParameters(
+            name=name,
+            dataset_id=dataset_id,
+            display_name=display_name,
+            metadata=metadata,
+            description=description,
+            encryption_spec=encryption_spec,
+            model_reference=model_reference,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _UpdateDatasetParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "datasets/{dataset_id}".format_map(request_url_dict)
+            else:
+                path = "datasets/{dataset_id}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "patch", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.Dataset._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _create_custom_job_resource(
+        self,
+        *,
+        custom_job: types.CustomJobOrDict,
+        config: Optional[types.VertexBaseConfigOrDict] = None,
+    ) -> types.CustomJob:
+        """
+        Creates a custom job.
+        """
+
+        parameter_model = types._CustomJobParameters(
+            custom_job=custom_job,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _CustomJobParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "customJobs".format_map(request_url_dict)
+            else:
+                path = "customJobs"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "post", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        if self._api_client.vertexai:
+            response_dict = _CustomJob_from_vertex(response_dict)
+
+        return_value = types.CustomJob._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _get_custom_job(
+        self, *, name: str, config: Optional[types.VertexBaseConfigOrDict] = None
+    ) -> types.CustomJob:
+        """
+        Gets a custom job.
+        """
+
+        parameter_model = types._GetCustomJobParameters(
+            name=name,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _GetCustomJobParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "customJobs/{name}".format_map(request_url_dict)
+            else:
+                path = "customJobs/{name}"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "get", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        if self._api_client.vertexai:
+            response_dict = _CustomJob_from_vertex(response_dict)
+
+        return_value = types.CustomJob._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    async def _optimize(
+        self,
+        *,
+        content: Optional[genai_types.ContentOrDict] = None,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> types.OptimizeResponseEndpoint:
+        """
+        Optimize a single prompt.
+        """
+
+        parameter_model = types._OptimizeRequestParameters(
+            content=content,
+            config=config,
+        )
+
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError(
+                "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode."
+            )
+        else:
+            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
+            else:
+                path = "tuningJobs:optimizePrompt"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "post", path, request_dict, http_options
+        )
+
+        response_dict = {} if not response.body else json.loads(response.body)
+
+        return_value = types.OptimizeResponseEndpoint._from_response(
+            response=response_dict,
+            kwargs=(
+                {
+                    "config": {
+                        "response_schema": getattr(
+                            parameter_model.config, "response_schema", None
+                        ),
+                        "response_json_schema": getattr(
+                            parameter_model.config, "response_json_schema", None
+                        ),
+                        "include_all_fields": getattr(
+                            parameter_model.config, "include_all_fields", None
+                        ),
+                    }
+                }
+                if getattr(parameter_model, "config", None)
+                else {}
+            ),
+        )
+
+        self._api_client._verify_response(return_value)
+        return return_value
+
+    _DEFAULT_TIMEOUT = 90
+    _DEFAULT_MAX_WAIT_TIME = 60
+
+    async def create(
+        self,
+        *,
+        prompt: types.PromptOrDict,
+        config: Optional[types.CreatePromptConfigOrDict] = None,
+    ) -> types.Prompt:
+        """Creates a new prompt in a Vertex Dataset resource.
+
+        This method waits for prompt creation to be complete before returning.
+
+        Note: This method does not create a versioned resource for your prompt.
+        Call create_version() to create a new prompt resource with a version.
+
+        Args:
+          prompt: The prompt to create.
+          config: Optional configuration for creating the prompt.
+
+        Returns:
+          A types.Prompt object representing the prompt with its associated
+          Dataset resources.
+        """
+        if isinstance(prompt, dict):
+            prompt = types.Prompt(**prompt)
+        if isinstance(config, dict):
+            config = types.CreatePromptConfig(**config)
+        elif not config:
+            config = types.CreatePromptConfig()
+
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
+
+        if prompt.prompt_data is None:
+            raise ValueError("Prompt data is required to create a prompt.")
+
+        prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
+            prompt,
+            variables=(
+                prompt.prompt_data.variables
+                if prompt.prompt_data and prompt.prompt_data.variables
+                else None
+            ),
+        )
+
+        # Step 1: Create the dataset resource for the prompt and wait for the operation to complete.
+        create_prompt_dataset_operation = await self._create_dataset_resource(
+            display_name=(
+                config.prompt_display_name
+                if config and config.prompt_display_name
+                else f"prompt_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}",
+            metadata_schema_uri=_prompt_management_utils.PROMPT_SCHEMA_URI,
+            metadata=prompt_metadata,
+            model_reference=prompt.prompt_data.model,
+            encryption_spec=(
+                config.encryption_spec if config and config.encryption_spec else None
+            ),
+        )
+        dataset_resource_name = await self._wait_for_operation(
+            operation=create_prompt_dataset_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_id = dataset_resource_name.split("/")[-1]
+
+        # Step 2: Get the dataset resource
+        dataset_resource = await self._get_dataset_resource(
+            name=dataset_id,
+        )
+        prompt._dataset = dataset_resource
+        return prompt
+
+    async def create_version(
+        self,
+        *,
+        prompt: types.PromptOrDict,
+        config: Optional[types.CreatePromptVersionConfigOrDict] = None,
+    ) -> types.Prompt:
+        """Creates a prompt resource and an initial prompt version.
+
+        When creating new prompt and prompt version resources, this waits for
+        the create operation to complete before returning.
+
+        Note: This method is recommended instead of create() since it creates a
+        versioned resource for your prompt.
+
+        Args:
+          prompt: The prompt to create.
+          config: Optional configuration for creating the prompt and prompt version.
+
+        Returns:
+          A types.Prompt object representing the prompt with its associated
+          Dataset and Dataset Version resources.
+        """
+        if isinstance(prompt, dict):
+            prompt = types.Prompt(**prompt)
+        if isinstance(config, dict):
+            config = types.CreatePromptVersionConfig(**config)
+        elif not config:
+            config = types.CreatePromptVersionConfig()
+
+        _prompt_management_utils._raise_for_invalid_prompt(prompt)
+
+        if prompt.prompt_data is None:
+            raise ValueError("Prompt data is required to create a prompt.")
+
+        prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
+            prompt,
+            variables=(
+                prompt.prompt_data.variables
+                if prompt.prompt_data and prompt.prompt_data.variables
+                else None
+            ),
+        )
+
+        # Step 1: Create the dataset resource for the prompt and wait for the operation to complete.
+        create_prompt_dataset_operation = await self._create_dataset_resource(
+            display_name=(
+                config.prompt_display_name
+                if config and config.prompt_display_name
+                else f"prompt_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}",
+            metadata_schema_uri=_prompt_management_utils.PROMPT_SCHEMA_URI,
+            metadata=prompt_metadata,
+            model_reference=prompt.prompt_data.model,
+            encryption_spec=(
+                config.encryption_spec if config and config.encryption_spec else None
+            ),
+        )
+        dataset_resource_name = await self._wait_for_operation(
+            operation=create_prompt_dataset_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_id = dataset_resource_name.split("/")[-1]
+
+        # Step 2: Get the dataset resource
+        dataset_resource = await self._get_dataset_resource(
+            name=dataset_id,
+        )
+        prompt._dataset = dataset_resource
+
+        if prompt._dataset.name is None:
+            raise ValueError("Invalid dataset resource.")
+
+        # Step 3: Create the dataset version
+        create_dataset_version_operation = await self._create_dataset_version_resource(
+            dataset_name=prompt._dataset.name.split("/")[-1],
+            display_name=(
+                config.version_display_name
+                if config and config.version_display_name is not None
+                else f"prompt_version_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+        )
+        dataset_version_resource_name = await self._wait_for_operation(
+            operation=create_dataset_version_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+
+        # Step 4: Get the dataset version resource and return it with the prompt
+        dataset_version_resource = await self._get_dataset_version_resource(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_resource_name.split("/")[-1],
+        )
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource
+        )
+        prompt._dataset = dataset_resource
+        prompt._dataset_version = dataset_version_resource
+        return prompt
+
+    async def update(
+        self,
+        *,
+        prompt_id: str,
+        prompt: types.PromptOrDict,
+        config: Optional[types.UpdatePromptConfigOrDict] = None,
+    ) -> types.Prompt:
+        """Updates an existing prompt and creates a new version for the prompt
+        associated with the provided prompt_id.
+
+        Args:
+          prompt_id: The ID of the prompt to create a new version for.
+          prompt: The updated prompt.
+          config: Optional configuration for updating the prompt.
+
+        Returns:
+          A types.Prompt object representing the updated prompt with its associated
+          Dataset and Dataset Version resources.
+        """
+
+        if isinstance(prompt, dict):
+            prompt = types.Prompt(**prompt)
+        if isinstance(config, dict):
+            config = types.UpdatePromptConfig(**config)
+        elif not config:
+            config = types.UpdatePromptConfig()
+
+        prompt_metadata = _prompt_management_utils._create_dataset_metadata_from_prompt(
+            prompt,
+            variables=(
+                prompt.prompt_data.variables
+                if prompt.prompt_data and prompt.prompt_data.variables
+                else None
+            ),
+        )
+
+        if not prompt.prompt_data:
+            raise ValueError("Prompt data is required to update a prompt.")
+
+        # Step 1: Update the dataset resource for the prompt and wait for the operation to complete.
+        updated_dataset_resource = await self._update_dataset_resource(
+            name=f"projects/{self._api_client.project}/locations/{self._api_client.location}/datasets/{prompt_id}",
+            dataset_id=prompt_id,
+            display_name=(
+                config.prompt_display_name
+                if config and config.prompt_display_name
+                else None
+            ),
+            metadata=prompt_metadata,
+            model_reference=prompt.prompt_data.model,
+            encryption_spec=(
+                config.encryption_spec if config and config.encryption_spec else None
+            ),
+            config=config,
+        )
+
+        if not updated_dataset_resource.name:
+            raise ValueError("Failed to update dataset resource.")
+
+        dataset_id = updated_dataset_resource.name.split("/")[-1]
+
+        # Step 2: Create a dataset version for the prompt.
+        create_dataset_version_operation = await self._create_dataset_version_resource(
+            dataset_name=dataset_id,
+            display_name=(
+                config.version_display_name
+                if config and config.version_display_name is not None
+                else f"prompt_version_{time.strftime('%Y%m%d-%H%M%S')}"
+            ),
+        )
+        dataset_version_resource_name = await self._wait_for_operation(
+            operation=create_dataset_version_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_version_id = dataset_version_resource_name.split("/")[-1]
+
+        # Step 3: Get the dataset version resource and return it with the prompt.
+        dataset_version_resource = await self._get_dataset_version_resource(
+            dataset_id=dataset_id,
+            dataset_version_id=dataset_version_id,
+        )
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource
+        )
+        prompt._dataset = updated_dataset_resource
+        prompt._dataset_version = dataset_version_resource
+        return prompt
+
+    async def _wait_for_operation(
+        self,
+        operation: types.DatasetOperation,
+        timeout: int,
+        max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
+    ) -> str:
+        """Waits for a dataset operation to complete.
+
+        Args:
+          operation: The dataset operation to wait for.
+          timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
+
+        Returns:
+          The name of the Dataset resource from the operation result.
+
+        Raises:
+          TimeoutError: If the operation does not complete within the timeout.
+          ValueError: If the operation fails.
+        """
+        done = False
+        prompt_dataset_operation: Optional[types.DatasetOperation] = None
+
+        response_operation_name = operation.name
+        if response_operation_name is None:
+            raise ValueError("Invalid operation name.")
+
+        dataset_id = response_operation_name.split("/datasets/")[1].split("/")[0]
+        operation_id = response_operation_name.split("/")[-1]
+
+        start_time = time.time()
+        sleep_duration = 5
+        wait_multiplier = 2
+        previous_time = time.time()
+
+        while not done:
+            if (time.time() - start_time) > timeout:
+                raise TimeoutError(
+                    "Create prompt operation did not complete within the"
+                    f" specified timeout of {timeout} seconds."
+                )
+            current_time = time.time()
+            if current_time - previous_time >= sleep_duration:
+                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
+                previous_time = current_time
+            await asyncio.sleep(sleep_duration)
+            prompt_dataset_operation = await self._get_dataset_operation(
+                dataset_id=dataset_id,
+                operation_id=operation_id,
+            )
+            done = (
+                (prompt_dataset_operation.done or False)
+                if hasattr(prompt_dataset_operation, "done")
+                else False
+            )
+        if (
+            not prompt_dataset_operation
+            or prompt_dataset_operation.response is None
+            or prompt_dataset_operation.response.get("name") is None
+        ):
+            raise ValueError("Error creating prompt version resource.")
+        if (
+            hasattr(prompt_dataset_operation, "error")
+            and prompt_dataset_operation.error is not None
+        ):
+            raise ValueError(
+                f"Error creating prompt version resource: {prompt_dataset_operation.error}"
+            )
+        return prompt_dataset_operation.response.get("name")  # type: ignore[return-value]
+
+    async def get(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.GetPromptConfig] = None,
+    ) -> types.Prompt:
+        """Gets a prompt resource from a Vertex Dataset.
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          config: Optional configuration for getting the prompt.
+
+        Returns:
+            A types.Prompt object representing the prompt with its associated Dataset and Dataset Version resources.
+        """
+
+        prompt_dataset_resource = await self._get_dataset_resource(name=prompt_id)
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            prompt_dataset_resource,
+        )
+        prompt._dataset = prompt_dataset_resource
+
+        return prompt
+
+    async def get_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.GetPromptConfig] = None,
+    ) -> types.Prompt:
+        """Gets a prompt resource from a Vertex Dataset.
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          version_id: The id of the Vertex Dataset Version resource containing the prompt version. For example, if the prompt version resource name is "projects/123/locations/us-central1/datasets/456/datasetVersions/1", then the version_id is "1".
+          config: Optional configuration for getting the prompt.
+
+        Returns:
+            A types.Prompt object representing the prompt with its associated Dataset and Dataset Version resources.
+        """
+
+        prompt_dataset_resource = await self._get_dataset_resource(name=prompt_id)
+        prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            prompt_dataset_resource,
+        )
+        prompt._dataset = prompt_dataset_resource
+
+        prompt_version_resource = await self._get_dataset_version_resource(
+            dataset_id=prompt_id,
+            dataset_version_id=version_id,
+        )
+        prompt._dataset_version = prompt_version_resource
+
+        return prompt
+
+    async def _wait_for_project_operation(
+        self,
+        operation: genai_types.ProjectOperation,
+        timeout: int,
+        max_wait_time: int = _DEFAULT_MAX_WAIT_TIME,
+    ) -> None:
+        """Waits for a dataset deletion operation to complete.
+
+        Delete operations are project level operations and are separate from dataset resource operations, for example: projects/123/locations/us-central1/operations/789.
+
+        Args:
+          operation: The project operation to wait for.
+          timeout: The maximum time to wait for the operation to complete.
+          max_wait_time: The maximum interval between polling requests in seconds.
+        Raises:
+          TimeoutError: If the operation does not complete within the timeout.
+          ValueError: If the operation fails.
+        """
+        done = False
+
+        start_time = time.time()
+        sleep_duration = 5
+        wait_multiplier = 2
+        previous_time = time.time()
+        while not done:
+            if (time.time() - start_time) > timeout:
+                raise TimeoutError(
+                    f"Delete operation did not complete within the"
+                    f" specified timeout of {timeout} seconds."
+                )
+            current_time = time.time()
+            if current_time - previous_time >= sleep_duration:
+                sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time)
+                previous_time = current_time
+            await asyncio.sleep(sleep_duration)
+            operations_module = operations.AsyncOperations(api_client_=self._api_client)
+
+            if operation.name is None:
+                raise ValueError("Invalid operation name.")
+            operation = await operations_module._get(
+                operation_id=operation.name.split("/")[-1],
+            )
+            done = (operation.done or False) if hasattr(operation, "done") else False
+        if hasattr(operation, "error") and operation.error is not None:
+            raise ValueError(f"Error in delete operation: {operation.error}")
+
+    async def delete(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.DeletePromptConfig] = None,
+    ) -> None:
+        """Deletes a prompt resource.
+
+        Args:
+          prompt_id: The id of the prompt resource to delete.
+
+        Raises:
+          TimeoutError: If the delete operation does not complete within the timeout.
+          ValueError: If the delete operation fails.
+        """
+
+        delete_prompt_operation = await self._delete_dataset(
+            prompt_id=prompt_id,
+            config=config,
+        )
+        await self._wait_for_project_operation(
+            operation=delete_prompt_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        logger.info(f"Deleted prompt with id: {prompt_id}")
+
+    async def delete_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.DeletePromptConfig] = None,
+    ) -> None:
+        """Deletes a prompt version resource.
+
+        Args:
+          prompt_id: The id of the prompt resource to delete.
+          version_id: The id of the prompt version resource to delete.
+
+        Raises:
+          TimeoutError: If the delete operation does not complete within the timeout.
+          ValueError: If the delete operation fails.
+        """
+        delete_version_operation = await self._delete_dataset_version(
+            prompt_id=prompt_id,
+            version_id=version_id,
+            config=config,
+        )
+
+        await self._wait_for_project_operation(
+            operation=delete_version_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        logger.info(
+            f"Deleted prompt version {version_id} from prompt with id: {prompt_id}"
+        )
+
+    async def _list_prompts_pager(
+        self,
+        *,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> AsyncPager[types.Dataset]:
+        return AsyncPager(
+            "datasets",
+            self._list_prompts,
+            await self._list_prompts(config=config),
+            config,
+        )
+
+    async def _list_versions_pager(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> AsyncPager[types.DatasetVersion]:
+        return AsyncPager(
+            "dataset_versions",
+            self._list_versions,
+            await self._list_versions(config=config, dataset_id=prompt_id),
+            config,
+        )
+
+    async def list(
+        self,
+        *,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> AsyncIterator[types.PromptRef]:
+        """Lists prompt resources in a project.
+
+        This method retrieves all the prompts from the project provided in the
+        agentplatform.Client constructor and returns a list of prompt references containing the prompt_id and model for the prompt.
+
+        To get the full types.Prompt resource for a PromptRef after calling this method, use the get() method with the prompt_id as the prompt_id argument.
+        Example usage:
+
+        ```
+        prompt_refs = client.aio.prompt_management.list_prompts()
+        async for prompt_ref in prompt_refs:
+          await client.prompt_management.get(prompt_id=prompt_ref.prompt_id)
+        ```
+
+        Args:
+          config: Optional configuration for listing prompts.
+
+        Returns:
+            An async iterator of types.PromptRef objects.
+        """
+        if isinstance(config, dict):
+            config = types.ListPromptsConfig(**config)
+        elif not config:
+            config = types.ListPromptsConfig()
+        async for dataset in await self._list_prompts_pager(config=config):
+            if not dataset or not dataset.model_reference or not dataset.name:
+                continue
+            prompt_ref = types.PromptRef(
+                model=dataset.model_reference, prompt_id=dataset.name.split("/")[-1]
+            )
+            yield prompt_ref
+
+    async def list_versions(
+        self,
+        *,
+        prompt_id: str,
+        config: Optional[types.ListPromptsConfigOrDict] = None,
+    ) -> AsyncIterator[types.PromptVersionRef]:
+        """Lists prompt version resources for a provided prompt_id.
+
+        This method retrieves all the prompt versions for a provided prompt_id.
+
+        To get the full types.Prompt resource for a PromptVersionRef after calling this method, use the get() method with the returned prompt_id and version_id.
+        Example usage:
+
+        ```
+        prompt_version_refs = await client.prompt_management.list_versions(prompt_id="123")
+        async for version_ref in prompt_version_refs:
+          await client.aio.prompt_management.get(prompt_id=version_ref.prompt_id, version_id=version_ref.version_id)
+        ```
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          config: Optional configuration for listing prompts.
+
+        Returns:
+            An async iterator of types.PromptVersionRef objects representing the prompt version resources for the provided prompt_id.
+
+        """
+        if isinstance(config, dict):
+            config = types.ListPromptsConfig(**config)
+        elif not config:
+            config = types.ListPromptsConfig()
+        async for dataset_version in await self._list_versions_pager(
+            config=config, prompt_id=prompt_id
+        ):
+            if (
+                not dataset_version
+                or not dataset_version.model_reference
+                or not dataset_version.name
+            ):
+                continue
+            prompt_version_ref = types.PromptVersionRef(
+                model=dataset_version.model_reference,
+                version_id=dataset_version.name.split("/")[-1],
+                prompt_id=prompt_id,
+            )
+            yield prompt_version_ref
+
+    async def restore_version(
+        self,
+        *,
+        prompt_id: str,
+        version_id: str,
+        config: Optional[types.RestoreVersionConfig] = None,
+    ) -> types.Prompt:
+        """Restores the provided prompt version to the latest version.
+
+        Args:
+          prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456".
+          version_id: The id of the Vertex Dataset Version resource to restore. For example, if the version resource name is "projects/123/locations/us-central1/datasets/456/datasetVersions/789", then the version_id is "789".
+          config: Optional configuration for restoring the prompt version.
+
+        Returns:
+            A types.Prompt object representing the prompt with the updated Dataset Version resource.
+        """
+
+        restore_prompt_operation = await self._restore_version(
+            dataset_id=prompt_id,
+            version_id=version_id,
+            config=config,
+        )
+        await self._wait_for_project_operation(
+            operation=restore_prompt_operation,
+            timeout=(
+                config.timeout
+                if config and config.timeout is not None
+                else self._DEFAULT_TIMEOUT
+            ),
+            max_wait_time=(
+                config.max_wait_time
+                if config and config.max_wait_time is not None
+                else self._DEFAULT_MAX_WAIT_TIME
+            ),
+        )
+        dataset_version_resource = await self._get_dataset_version_resource(
+            dataset_id=prompt_id,
+            dataset_version_id=version_id,
+        )
+        updated_prompt = _prompt_management_utils._create_prompt_from_dataset_metadata(
+            dataset_version_resource,
+        )
+        updated_prompt._dataset_version = dataset_version_resource
+        return updated_prompt
+
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI prompts.launch_optimization_job method is "
+        "experimental, and may change in future versions."
+    )
+    async def launch_optimization_job(
+        self,
+        method: types.PromptOptimizerMethod,
+        config: types.PromptOptimizerConfigOrDict,
+    ) -> types.CustomJob:
+        """Call async Vertex AI Prompt Optimizer (VAPO).
+
+
+        Note: The `wait_for_completion` parameter in the config will be
+        ignored when using the AsyncClient, as it is not supported.
+
+        Example usage:
+        client = agentplatform.Client(project=PROJECT_NAME, location='us-central1')
+        vapo_config = agentplatform.types.PromptOptimizerConfig(
+            config_path='gs://you-bucket-name/your-config.json',
+            service_account=service_account,
+        )
+        job = await client.aio.prompts.launch_optimization_job(
+            method=types.PromptOptimizerMethod.VAPO, config=vapo_config)
+
+        Args:
+          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 isinstance(config, dict):
+            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/vaipo:preview_android_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(
+                "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:
+            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+            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])
+
+        region = self._api_client.location
+        project = self._api_client.project
+        container_args = {
+            "config": config.config_path,
+        }
+        args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
+        worker_pool_specs = [
+            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)
+
+        job_spec = types.CustomJobSpec(
+            worker_pool_specs=worker_pool_specs,
+            base_output_directory=genai_types.GcsDestination(output_uri_prefix=bucket),
+            service_account=service_account,
+        )
+
+        custom_job = types.CustomJob(
+            display_name=display_name,
+            job_spec=job_spec,
+        )
+
+        job = await self._create_custom_job_resource(
+            custom_job=custom_job,
+        )
+
+        # 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)
+
+        # Construct the dashboard URL to show to the user.
+        dashboard_url = f"https://console.cloud.google.com/vertex-ai/locations/{region}/training/{job_id}/cpu?project={project}"
+        logger.info("View the job status at: %s", dashboard_url)
+
+        return job
+
+    async def _custom_optimize(
+        self,
+        *,
+        content: Optional[genai_types.ContentOrDict] = None,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> types.OptimizeResponse:
+        """Optimize a single prompt."""
+        if isinstance(config, dict):
+            config.pop("examples_dataframe", None)
+        elif config and hasattr(config, "examples_dataframe"):
+            del config.examples_dataframe
+
+        parameter_model = types._OptimizeRequestParameters(
+            content=content,
+            config=config,
+        )
+        request_url_dict: Optional[dict[str, str]]
+        if not self._api_client.vertexai:
+            raise ValueError("This method is only supported in the Vertex AI client.")
+        else:
+            request_dict = _OptimizeRequestParameters_to_vertex(parameter_model)
+            request_url_dict = request_dict.get("_url")
+            if request_url_dict:
+                path = "tuningJobs:optimizePrompt".format_map(request_url_dict)
+            else:
+                path = "tuningJobs:optimizePrompt"
+
+        query_params = request_dict.get("_query")
+        if query_params:
+            path = f"{path}?{urlencode(query_params)}"
+        # TODO: remove the hack that pops config.
+        request_dict.pop("config", None)
+
+        http_options: Optional[types.HttpOptions] = None
+        if (
+            parameter_model.config is not None
+            and parameter_model.config.http_options is not None
+        ):
+            http_options = parameter_model.config.http_options
+
+        request_dict = _common.convert_to_dict(request_dict)
+        request_dict = _common.encode_unserializable_types(request_dict)
+
+        response = await self._api_client.async_request(
+            "post", path, request_dict, http_options
+        )
+
+        response_list = "" if not response.body else json.loads(response.body)
+
+        return_value = []
+
+        for response_dict in response_list:
+            response_value = types.OptimizeResponseEndpoint._from_response(
+                response=response_dict, kwargs=parameter_model.model_dump()
+            )
+            self._api_client._verify_response(response_value)
+            content = response_value.content
+            if content is not None:
+                parts = content.parts
+                if parts and parts[0].text is not None:
+                    return_value.append(parts[0].text)
+
+        output = "".join(return_value)
+        final_response = types.OptimizeResponse(raw_text_response=output)
+        try:
+            final_response.parsed_response = _prompt_optimizer_utils._parse(output)
+        except Exception as e:
+            logger.warning(
+                f"Failed to parse response: {e}. Returning only raw_text_response."
+            )
+        return final_response
+
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI prompts.optimize method is "
+        "experimental, and may change in future versions."
+    )
+    async def optimize(
+        self,
+        *,
+        prompt: str,
+        config: Optional[types.OptimizeConfigOrDict] = None,
+    ) -> types.OptimizeResponse:
+        """Makes an async request to the optimizePrompt endpoint and returns an optimized prompt.
+
+        Example usage:
+        client = agentplatform.Client(project=PROJECT_NAME, location='us-central1')
+        prompt = "Generate system instructions for analyzing medical articles"
+        response = await client.aio.prompts.optimize(prompt=prompt)
+
+        Args:
+          prompt: Required. The prompt to optimize.
+          config: Optional.The configuration for prompt optimization. To optimize
+            prompts from Android API provide
+            types.OptimizeConfig(
+                optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO
+            )
+            For few-shot optimization, provide:
+            optim_target = types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS # or types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
+            types.OptimizeConfig(
+                optimization_target=optim_target,
+                examples_dataframe=dataframe
+            )
+            OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS indicates that the few-shot
+              examples include specific scoring rubrics and their corresponding
+              evaluations.
+            OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE indicates that the few-shot
+              examples include a ground-truth target response.
+        Returns:
+          The parsed response from the API request.
+        """
+        if isinstance(config, dict):
+            config = types.OptimizeConfig(**config)
+
+        optimization_target: Optional[types.OptimizeTarget] = None
+        if config is not None:
+            optimization_target = config.optimization_target
+
+        final_prompt = prompt
+        if (
+            optimization_target
+            == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
+            or optimization_target
+            == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
+        ):
+            final_prompt = _prompt_optimizer_utils._get_few_shot_prompt(prompt, config)
+
+        # TODO: b/435653980 - replace the custom method with a generated method.
+        config_for_api = config.model_copy() if config else None
+        return await self._custom_optimize(
+            content=genai_types.Content(
+                parts=[genai_types.Part(text=final_prompt)], role="user"
+            ),
+            config=config_for_api,
+        )