datasets.py

# 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 json
import logging
import time
from typing import Any, Optional, Union
from urllib.parse import urlencode

from google.genai import _api_module
from google.genai import _common
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
import pandas as pd

from . import _datasets_utils
from . import types

logger = logging.getLogger("vertexai_genai.datasets")


def _AssembleDatasetParameters_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, ["gemini_request_read_config"]) is not None:
        setv(
            to_object,
            ["geminiRequestReadConfig"],
            getv(from_object, ["gemini_request_read_config"]),
        )

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _AssessDatasetParameters_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, ["gemini_request_read_config"]) is not None:
        setv(
            to_object,
            ["geminiRequestReadConfig"],
            getv(from_object, ["gemini_request_read_config"]),
        )

    if getv(from_object, ["tuning_resource_usage_assessment_config"]) is not None:
        setv(
            to_object,
            ["tuningResourceUsageAssessmentConfig"],
            getv(from_object, ["tuning_resource_usage_assessment_config"]),
        )

    if getv(from_object, ["tuning_validation_assessment_config"]) is not None:
        setv(
            to_object,
            ["tuningValidationAssessmentConfig"],
            getv(from_object, ["tuning_validation_assessment_config"]),
        )

    if (
        getv(from_object, ["batch_prediction_resource_usage_assessment_config"])
        is not None
    ):
        setv(
            to_object,
            ["batchPredictionResourceUsageAssessmentConfig"],
            getv(from_object, ["batch_prediction_resource_usage_assessment_config"]),
        )

    if getv(from_object, ["batch_prediction_validation_assessment_config"]) is not None:
        setv(
            to_object,
            ["batchPredictionValidationAssessmentConfig"],
            getv(from_object, ["batch_prediction_validation_assessment_config"]),
        )

    if getv(from_object, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _CreateMultimodalDatasetParameters_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, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


def _DeleteMultimodalDatasetRequestParameters_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 _GetMultimodalDatasetOperationParameters_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 _GetMultimodalDatasetParameters_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 _ListMultimodalDatasetsConfig_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 _ListMultimodalDatasetsRequestParameters_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"],
            _ListMultimodalDatasetsConfig_to_vertex(
                getv(from_object, ["config"]), to_object
            ),
        )

    return to_object


def _UpdateMultimodalDatasetParameters_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, ["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, ["config"]) is not None:
        setv(to_object, ["config"], getv(from_object, ["config"]))

    return to_object


class Datasets(_api_module.BaseModule):

    def _assemble_multimodal_dataset(
        self,
        *,
        name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssembleDatasetConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Assembles a multimodal dataset resource.
        """

        parameter_model = types._AssembleDatasetParameters(
            name=name,
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _AssembleDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}:assemble".format_map(request_url_dict)
            else:
                path = "{name}:assemble"

        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.MultimodalDatasetOperation._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 _assess_multimodal_dataset(
        self,
        *,
        name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        tuning_resource_usage_assessment_config: Optional[
            types.TuningResourceUsageAssessmentConfigOrDict
        ] = None,
        tuning_validation_assessment_config: Optional[
            types.TuningValidationAssessmentConfigOrDict
        ] = None,
        batch_prediction_resource_usage_assessment_config: Optional[
            types.BatchPredictionResourceUsageAssessmentConfigOrDict
        ] = None,
        batch_prediction_validation_assessment_config: Optional[
            types.BatchPredictionValidationAssessmentConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Assesses a multimodal dataset resource.
        """

        parameter_model = types._AssessDatasetParameters(
            name=name,
            gemini_request_read_config=gemini_request_read_config,
            tuning_resource_usage_assessment_config=tuning_resource_usage_assessment_config,
            tuning_validation_assessment_config=tuning_validation_assessment_config,
            batch_prediction_resource_usage_assessment_config=batch_prediction_resource_usage_assessment_config,
            batch_prediction_validation_assessment_config=batch_prediction_validation_assessment_config,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _AssessDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}:assess".format_map(request_url_dict)
            else:
                path = "{name}:assess"

        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.MultimodalDatasetOperation._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_multimodal_dataset(
        self,
        *,
        name: Optional[str] = None,
        display_name: Optional[str] = None,
        metadata_schema_uri: Optional[str] = None,
        metadata: Optional[types.SchemaTablesDatasetMetadataOrDict] = None,
        description: Optional[str] = None,
        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Creates a dataset resource to store multimodal datasets.
        """

        parameter_model = types._CreateMultimodalDatasetParameters(
            name=name,
            display_name=display_name,
            metadata_schema_uri=metadata_schema_uri,
            metadata=metadata,
            description=description,
            encryption_spec=encryption_spec,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _CreateMultimodalDatasetParameters_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.MultimodalDatasetOperation._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_multimodal_dataset(
        self, *, name: str, config: Optional[types.VertexBaseConfigOrDict] = None
    ) -> types.MultimodalDatasetOperation:
        """
        Deletes a multimodal dataset resource.
        """

        parameter_model = types._DeleteMultimodalDatasetRequestParameters(
            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 the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _DeleteMultimodalDatasetRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{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("delete", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.MultimodalDatasetOperation._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_multimodal_dataset(
        self,
        *,
        name: Optional[str] = None,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """
        Gets a multimodal dataset resource.
        """

        parameter_model = types._GetMultimodalDatasetParameters(
            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 the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _GetMultimodalDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{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.MultimodalDataset._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_multimodal_dataset_operation(
        self,
        *,
        dataset_id: Optional[str] = None,
        operation_id: Optional[str] = None,
        config: Optional[types.GetMultimodalDatasetOperationConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Gets the operation from creating a multimodal dataset.
        """

        parameter_model = types._GetMultimodalDatasetOperationParameters(
            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 the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _GetMultimodalDatasetOperationParameters_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.MultimodalDatasetOperation._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_multimodal_datasets(
        self, *, config: Optional[types.ListMultimodalDatasetsConfigOrDict] = None
    ) -> types.ListMultimodalDatasetsResponse:
        """
        Lists multimodal datasets.
        """

        parameter_model = types._ListMultimodalDatasetsRequestParameters(
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _ListMultimodalDatasetsRequestParameters_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.ListMultimodalDatasetsResponse._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_multimodal_dataset(
        self,
        *,
        name: Optional[str] = None,
        display_name: Optional[str] = None,
        metadata: Optional[types.SchemaTablesDatasetMetadataOrDict] = None,
        description: Optional[str] = None,
        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """
        Updates a multimodal dataset resource.
        """

        parameter_model = types._UpdateMultimodalDatasetParameters(
            name=name,
            display_name=display_name,
            metadata=metadata,
            description=description,
            encryption_spec=encryption_spec,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _UpdateMultimodalDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{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("patch", path, request_dict, http_options)

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.MultimodalDataset._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 _wait_for_operation(
        self,
        operation: types.MultimodalDatasetOperation,
        timeout_seconds: int,
    ) -> dict[str, Any]:
        """Waits for a multimodal or assemble dataset operation to complete.

        Args:
          operation: The multimodal or assemble dataset operation to wait for.
          timeout_seconds: The maximum time in seconds to wait for the operation
          to complete.

        Returns:
            A dict containing the operation response.

        Raises:
          TimeoutError: If the operation does not complete within the timeout.
          ValueError: If the operation fails.
        """
        response_operation_name = operation.name
        if response_operation_name is None:
            raise ValueError("Dataset operation name is empty.")
        dataset_id = response_operation_name.split("/datasets/")[1].split("/")[0]
        operation_id = response_operation_name.split("/")[-1]

        start_time = time.time()
        sleep_duration_seconds = 5
        wait_multiplier = 2
        max_wait_time_seconds = 60

        while (time.time() - start_time) < timeout_seconds:
            operation = self._get_multimodal_dataset_operation(
                dataset_id=dataset_id,
                operation_id=operation_id,
            )
            if operation.done:
                break
            time.sleep(sleep_duration_seconds)
            sleep_duration_seconds = min(
                sleep_duration_seconds * wait_multiplier, max_wait_time_seconds
            )
        else:
            raise TimeoutError(
                "The operation did not complete within the"
                f" specified timeout of {timeout_seconds} seconds."
            )
        if not operation or operation.response is None:
            logger.error(f"Error running the operation {operation.response}.")
            raise ValueError(f"Error running the operation {operation.response}.")
        if hasattr(operation, "error") and operation.error is not None:
            raise ValueError(f"Error running the operation {operation.error}")
        return operation.response

    def create_from_bigquery(
        self,
        *,
        multimodal_dataset: types.MultimodalDatasetOrDict,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Creates a multimodal dataset from a BigQuery table.

        Args:
          multimodal_dataset:
            Required. A representation of a multimodal dataset.
          config:
            Optional. A configuration for creating the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDataset object representing a multimodal dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        _datasets_utils.validate_multimodal_dataset_bigquery_uri(multimodal_dataset)

        if isinstance(config, dict):
            config = types.CreateMultimodalDatasetConfig(**config)
        elif not config:
            config = types.CreateMultimodalDatasetConfig()

        display_name = (
            multimodal_dataset.display_name
            if multimodal_dataset.display_name is not None
            else _datasets_utils.generate_multimodal_dataset_display_name()
        )
        multimodal_dataset_operation = self._create_multimodal_dataset(
            config=config,
            display_name=display_name,
            metadata_schema_uri=_datasets_utils.METADATA_SCHEMA_URI,
            metadata=multimodal_dataset.metadata,
        )
        response = self._wait_for_operation(
            operation=multimodal_dataset_operation,
            timeout_seconds=config.timeout,
        )
        return _datasets_utils.create_from_response(
            types.MultimodalDataset, response, config
        )

    def create_from_pandas(
        self,
        *,
        dataframe: pd.DataFrame,
        multimodal_dataset: Optional[types.MultimodalDatasetOrDict] = None,
        target_table_id: Optional[str] = None,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Creates a multimodal dataset from a pandas dataframe.

        Args:
            dataframe (pandas.DataFrame):
                The pandas dataframe to be used for the created dataset.
            multimodal_dataset:
                Optional. A representation of a multimodal dataset.
            target_table_id (str):
                Optional. The BigQuery table id where the dataframe will be
                uploaded. The table id can be in the format of "dataset.table"
                or "project.dataset.table". Note that the BigQuery
                dataset must already exist and be in the same location as the
                multimodal dataset. If not provided, a generated table id will
                be created in the `vertex_datasets` dataset (e.g.
                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
            config:
                Optional. A configuration for creating the multimodal dataset. If not
                provided, the default configuration will be used.

        Returns:
            dataset (MultimodalDataset):
                The created multimodal dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        elif not multimodal_dataset:
            multimodal_dataset = types.MultimodalDataset()

        bigframes = _datasets_utils._try_import_bigframes()
        project = self._api_client.project
        location = self._api_client.location
        credentials = self._api_client._credentials

        session_options = bigframes.BigQueryOptions(
            credentials=credentials,
            project=project,
            location=location,
        )
        with bigframes.connect(session_options) as session:
            return self.create_from_bigframes(
                dataframe=session.read_pandas(dataframe),
                multimodal_dataset=multimodal_dataset,
                target_table_id=target_table_id,
                config=config,
            )

    def create_from_bigframes(
        self,
        *,
        dataframe: "bigframes.pandas.DataFrame",  # type: ignore # noqa: F821
        multimodal_dataset: Optional[types.MultimodalDatasetOrDict] = None,
        target_table_id: Optional[str] = None,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Creates a multimodal dataset from a bigframes dataframe.

        Args:
            dataframe (bigframes.pandas.DataFrame):
                The BigFrames dataframe that will be used for the created
                dataset.
            multimodal_dataset:
                Optional. A representation of a multimodal dataset.
            target_table_id (str):
                Optional. The BigQuery table id where the dataframe will be
                uploaded. The table id can be in the format of "dataset.table"
                or "project.dataset.table". Note that the BigQuery
                dataset must already exist and be in the same location as the
                multimodal dataset. If not provided, a generated table id will
                be created in the `vertex_datasets` dataset (e.g.
                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
            config:
                Optional. A configuration for creating the multimodal dataset. If not
                provided, the default configuration will be used.

        Returns:
            dataset (MultimodalDataset):
                The created multimodal dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        elif not multimodal_dataset:
            multimodal_dataset = types.MultimodalDataset()

        bigquery = _datasets_utils._try_import_bigquery()
        project = self._api_client.project
        location = self._api_client.location
        credentials = self._api_client._credentials

        if target_table_id:
            target_table_id = _datasets_utils._normalize_and_validate_table_id(
                table_id=target_table_id,
                project=project,
                location=location,
                credentials=credentials,
            )
        else:
            dataset_id = _datasets_utils._create_default_bigquery_dataset_if_not_exists(
                project=project, location=location, credentials=credentials
            )
            target_table_id = _datasets_utils._generate_target_table_id(dataset_id)

        client = bigquery.Client(project=project, credentials=credentials)
        _datasets_utils.save_dataframe_to_bigquery(
            dataframe,
            target_table_id,
            client,
        )

        return self.create_from_bigquery(
            multimodal_dataset=multimodal_dataset.model_copy(
                update={
                    "metadata": types.SchemaTablesDatasetMetadata(
                        input_config=types.SchemaTablesDatasetMetadataInputConfig(
                            bigquery_source=types.SchemaTablesDatasetMetadataBigQuerySource(
                                uri=f"bq://{target_table_id}"
                            )
                        )
                    )
                }
            ),
            config=config,
        )

    def update_multimodal_dataset(
        self,
        *,
        multimodal_dataset: types.MultimodalDatasetOrDict,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Updates a multimodal dataset.

        Updatable fields include:
        - display_name
        - description

        Args:
          multimodal_dataset:
            Required. A representation of a multimodal dataset.
          config:
            Optional. A configuration for updating the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDataset object representing the retrieved multimodal
          dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        _datasets_utils.validate_multimodal_dataset_bigquery_uri(multimodal_dataset)

        if isinstance(config, dict):
            config = types.VertexBaseConfig(**config)
        elif not config:
            config = types.VertexBaseConfig()

        return self._update_multimodal_dataset(
            config=config,
            name=multimodal_dataset.name,
            display_name=multimodal_dataset.display_name,
            description=multimodal_dataset.description,
            metadata=multimodal_dataset.metadata,
        )

    def get_multimodal_dataset(
        self,
        *,
        name: str,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Gets a multimodal dataset.

        Args:
          name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          config:
            Optional. A configuration for getting the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDataset object representing the retrieved multimodal
          dataset.
        """
        if isinstance(config, dict):
            config = types.VertexBaseConfig(**config)
        elif not config:
            config = types.VertexBaseConfig()

        name = _datasets_utils.resolve_dataset_name(
            name, self._api_client.project, self._api_client.location
        )

        return self._get_multimodal_dataset(config=config, name=name)

    def delete_multimodal_dataset(
        self,
        *,
        name: str,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """Deletes a multimodal dataset.

        Args:
          name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          config:
            Optional. A configuration for deleting the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDatasetOperation object representing the delete
          multimodal dataset operation.
        """
        if isinstance(config, dict):
            config = types.VertexBaseConfig(**config)
        elif not config:
            config = types.VertexBaseConfig()

        name = _datasets_utils.resolve_dataset_name(
            name, self._api_client.project, self._api_client.location
        )

        return self._delete_multimodal_dataset(config=config, name=name)

    def assemble(
        self,
        *,
        name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssembleDatasetConfigOrDict] = None,
    ) -> str:
        """Assemble the dataset into a BigQuery table.

        Waits for the assemble operation to complete before returning.

        Args:
          name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          gemini_request_read_config:
            Optional. The read config to use to assemble the dataset. If
            not provided, the read config attached to the dataset will be
            used.
          config:
            Optional. A configuration for assembling the dataset. If not
            provided, the default configuration will be used.

        Returns:
            The URI of the bigquery table of the assembled dataset.
        """
        if isinstance(config, dict):
            config = types.AssembleDatasetConfig(**config)
        elif not config:
            config = types.AssembleDatasetConfig()

        name = _datasets_utils.resolve_dataset_name(
            name, self._api_client.project, self._api_client.location
        )

        operation = self._assemble_multimodal_dataset(
            name=name,
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        return response["bigqueryDestination"]  # type: ignore[no-any-return]

    def assess_tuning_resources(
        self,
        *,
        dataset_name: str,
        model_name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.TuningResourceUsageAssessmentResult:
        """Assess the tuning resources required for a given model.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
            Required. The name of the model to assess the tuning resources
            for.
          gemini_request_read_config:
            Optional. The read config used to assemble the dataset
            before assessing the tuning resources. If not provided, the
            read config attached to the dataset will be used. Required
            if no read config is attached to the dataset.
          config:
            Optional. A configuration for assessing the tuning resources. If not
            provided, the default configuration will be used.

        Returns:
          A types.TuningResourceUsageAssessmentResult object representing the
          tuning resource usage assessment result.
        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = self._assess_multimodal_dataset(
            name=dataset_name,
            tuning_resource_usage_assessment_config=types.TuningResourceUsageAssessmentConfig(
                model_name=model_name
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        return _datasets_utils.create_from_response(
            types.TuningResourceUsageAssessmentResult,
            response["tuningResourceUsageAssessmentResult"],
            config,
        )

    def assess_tuning_validity(
        self,
        *,
        dataset_name: str,
        model_name: str,
        dataset_usage: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.TuningValidationAssessmentResult:
        """Assess if the assembled dataset is valid in terms of tuning a given
        model.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
              Required. The name of the model to assess the tuning validity
              for.
          dataset_usage:
              Required. The dataset usage to assess the tuning validity for.
              Must be one of the following: SFT_TRAINING, SFT_VALIDATION.
          gemini_request_read_config:
              Optional. The read config used to assemble the dataset
              before assessing the tuning validity. If not provided, the
              read config attached to the dataset will be used. Required
              if no read config is attached to the dataset.
          config:
            Optional. A configuration for assessing the tuning validity. If not
            provided, the default configuration will be used.

        Returns:
            A dict containing the tuning validity assessment result. The dict
            contains the following keys:
            - errors: A list of errors that occurred during the tuning validity
              assessment.
        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = self._assess_multimodal_dataset(
            name=dataset_name,
            tuning_validation_assessment_config=types.TuningValidationAssessmentConfig(
                model_name=model_name,
                dataset_usage=dataset_usage,
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        return _datasets_utils.create_from_response(
            types.TuningValidationAssessmentResult,
            response["tuningValidationAssessmentResult"],
            config,
        )

    def assess_batch_prediction_resources(
        self,
        *,
        dataset_name: str,
        model_name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.BatchPredictionResourceUsageAssessmentResult:
        """Assess the batch prediction resources required for a given model.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
              Required. The name of the model to assess the batch prediction
              resources.
          gemini_request_read_config:
              Optional. The read config used to assemble the dataset
              before assessing the batch prediction resources. If not provided,
              the read config attached to the dataset will be used. Required
              if no read config is attached to the dataset.
          config:
              Optional. A configuration for assessing the batch prediction
              resources. If not provided, the default configuration will be
              used.

        Returns:
            A types.BatchPredictionResourceUsageAssessmentResult object
            representing the batch prediction resource usage assessment result.
            It contains the following keys:
              - token_count: The number of tokens in the dataset.
              - audio_token_count: The number of audio tokens in the dataset.

        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = self._assess_multimodal_dataset(
            name=dataset_name,
            batch_prediction_resource_usage_assessment_config=types.BatchPredictionResourceUsageAssessmentConfig(
                model_name=model_name,
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        result = response["batchPredictionResourceUsageAssessmentResult"]
        return _datasets_utils.create_from_response(
            types.BatchPredictionResourceUsageAssessmentResult, result, config
        )

    def assess_batch_prediction_validity(
        self,
        *,
        dataset_name: str,
        model_name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.BatchPredictionValidationAssessmentResult:
        """Assess if the assembled dataset is valid in terms of batch prediction
          for a given model. Raises an error if the dataset is invalid, otherwise
          returns None.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
            Required. The name of the model to assess the batch prediction
            validity for.
          gemini_request_read_config:
              Optional. The read config used to assemble the dataset
              before assessing the batch prediction validity. If not provided, the
              read config attached to the dataset will be used. Required
              if no read config is attached to the dataset.
          config:
            Optional. A configuration for assessing the batch prediction validity.
            If not provided, the default configuration will be used.

        Returns:
            A types.BatchPredictionValidationAssessmentResult object representing
            the batch prediction validity assessment result.
            It contains the following keys:
              - errors: A list of errors that occurred during the batch prediction
                validity assessment.
        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = self._assess_multimodal_dataset(
            name=dataset_name,
            batch_prediction_validation_assessment_config=types.BatchPredictionValidationAssessmentConfig(
                model_name=model_name,
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        result = response["batchPredictionValidationAssessmentResult"]
        return _datasets_utils.create_from_response(
            types.BatchPredictionValidationAssessmentResult, result, config
        )


class AsyncDatasets(_api_module.BaseModule):

    async def _assemble_multimodal_dataset(
        self,
        *,
        name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssembleDatasetConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Assembles a multimodal dataset resource.
        """

        parameter_model = types._AssembleDatasetParameters(
            name=name,
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _AssembleDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}:assemble".format_map(request_url_dict)
            else:
                path = "{name}:assemble"

        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.MultimodalDatasetOperation._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 _assess_multimodal_dataset(
        self,
        *,
        name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        tuning_resource_usage_assessment_config: Optional[
            types.TuningResourceUsageAssessmentConfigOrDict
        ] = None,
        tuning_validation_assessment_config: Optional[
            types.TuningValidationAssessmentConfigOrDict
        ] = None,
        batch_prediction_resource_usage_assessment_config: Optional[
            types.BatchPredictionResourceUsageAssessmentConfigOrDict
        ] = None,
        batch_prediction_validation_assessment_config: Optional[
            types.BatchPredictionValidationAssessmentConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Assesses a multimodal dataset resource.
        """

        parameter_model = types._AssessDatasetParameters(
            name=name,
            gemini_request_read_config=gemini_request_read_config,
            tuning_resource_usage_assessment_config=tuning_resource_usage_assessment_config,
            tuning_validation_assessment_config=tuning_validation_assessment_config,
            batch_prediction_resource_usage_assessment_config=batch_prediction_resource_usage_assessment_config,
            batch_prediction_validation_assessment_config=batch_prediction_validation_assessment_config,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _AssessDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}:assess".format_map(request_url_dict)
            else:
                path = "{name}:assess"

        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.MultimodalDatasetOperation._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_multimodal_dataset(
        self,
        *,
        name: Optional[str] = None,
        display_name: Optional[str] = None,
        metadata_schema_uri: Optional[str] = None,
        metadata: Optional[types.SchemaTablesDatasetMetadataOrDict] = None,
        description: Optional[str] = None,
        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Creates a dataset resource to store multimodal datasets.
        """

        parameter_model = types._CreateMultimodalDatasetParameters(
            name=name,
            display_name=display_name,
            metadata_schema_uri=metadata_schema_uri,
            metadata=metadata,
            description=description,
            encryption_spec=encryption_spec,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _CreateMultimodalDatasetParameters_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.MultimodalDatasetOperation._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_multimodal_dataset(
        self, *, name: str, config: Optional[types.VertexBaseConfigOrDict] = None
    ) -> types.MultimodalDatasetOperation:
        """
        Deletes a multimodal dataset resource.
        """

        parameter_model = types._DeleteMultimodalDatasetRequestParameters(
            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 the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _DeleteMultimodalDatasetRequestParameters_to_vertex(
                parameter_model
            )
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{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(
            "delete", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.MultimodalDatasetOperation._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_multimodal_dataset(
        self,
        *,
        name: Optional[str] = None,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """
        Gets a multimodal dataset resource.
        """

        parameter_model = types._GetMultimodalDatasetParameters(
            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 the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _GetMultimodalDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{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.MultimodalDataset._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_multimodal_dataset_operation(
        self,
        *,
        dataset_id: Optional[str] = None,
        operation_id: Optional[str] = None,
        config: Optional[types.GetMultimodalDatasetOperationConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """
        Gets the operation from creating a multimodal dataset.
        """

        parameter_model = types._GetMultimodalDatasetOperationParameters(
            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 the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _GetMultimodalDatasetOperationParameters_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.MultimodalDatasetOperation._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_multimodal_datasets(
        self, *, config: Optional[types.ListMultimodalDatasetsConfigOrDict] = None
    ) -> types.ListMultimodalDatasetsResponse:
        """
        Lists multimodal datasets.
        """

        parameter_model = types._ListMultimodalDatasetsRequestParameters(
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _ListMultimodalDatasetsRequestParameters_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.ListMultimodalDatasetsResponse._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_multimodal_dataset(
        self,
        *,
        name: Optional[str] = None,
        display_name: Optional[str] = None,
        metadata: Optional[types.SchemaTablesDatasetMetadataOrDict] = None,
        description: Optional[str] = None,
        encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """
        Updates a multimodal dataset resource.
        """

        parameter_model = types._UpdateMultimodalDatasetParameters(
            name=name,
            display_name=display_name,
            metadata=metadata,
            description=description,
            encryption_spec=encryption_spec,
            config=config,
        )

        request_url_dict: Optional[dict[str, str]]
        if not self._api_client.vertexai:
            raise ValueError(
                "This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
            )
        else:
            request_dict = _UpdateMultimodalDatasetParameters_to_vertex(parameter_model)
            request_url_dict = request_dict.get("_url")
            if request_url_dict:
                path = "{name}".format_map(request_url_dict)
            else:
                path = "{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(
            "patch", path, request_dict, http_options
        )

        response_dict = {} if not response.body else json.loads(response.body)

        return_value = types.MultimodalDataset._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 _wait_for_operation(
        self,
        operation: types.MultimodalDatasetOperation,
        timeout_seconds: int,
    ) -> dict[str, Any]:
        """Waits for a multimodal dataset operation to complete.

        Args:
          operation: The multimodal dataset operation to wait for.
          timeout_seconds: The maximum time in seconds to wait for the operation
          to complete.

        Returns:
            A dict containing the operation response.

        Raises:
          TimeoutError: If the operation does not complete within the timeout.
          ValueError: If the operation fails.
        """
        response_operation_name = operation.name
        if response_operation_name is None:
            raise ValueError("Dataset operation name is empty.")
        dataset_id = response_operation_name.split("/datasets/")[1].split("/")[0]
        operation_id = response_operation_name.split("/")[-1]

        start_time = time.time()
        sleep_duration_seconds = 5
        wait_multiplier = 2
        max_wait_time_seconds = 60

        while (time.time() - start_time) < timeout_seconds:
            operation = await self._get_multimodal_dataset_operation(
                dataset_id=dataset_id,
                operation_id=operation_id,
            )
            if operation.done:
                break
            await asyncio.sleep(sleep_duration_seconds)
            sleep_duration_seconds = min(
                sleep_duration_seconds * wait_multiplier, max_wait_time_seconds
            )
        else:
            raise TimeoutError(
                "The operation did not complete within the"
                f" specified timeout of {timeout_seconds} seconds."
            )
        if not operation or operation.response is None:
            logger.error(f"Error running the operation {operation.response}.")
            raise ValueError(f"Error running the operation {operation.response}.")
        if hasattr(operation, "error") and operation.error is not None:
            raise ValueError(f"Error running the operation {operation.error}")
        return operation.response

    async def create_from_bigquery(
        self,
        *,
        multimodal_dataset: types.MultimodalDatasetOrDict,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Creates a multimodal dataset from a BigQuery table.

        Args:
          multimodal_dataset:
            Required. A representation of a multimodal dataset.
          config:
            Optional. A configuration for creating the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDataset object representing a multimodal dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        _datasets_utils.validate_multimodal_dataset_bigquery_uri(multimodal_dataset)

        if isinstance(config, dict):
            config = types.CreateMultimodalDatasetConfig(**config)
        elif not config:
            config = types.CreateMultimodalDatasetConfig()

        display_name = (
            multimodal_dataset.display_name
            if multimodal_dataset.display_name is not None
            else _datasets_utils.generate_multimodal_dataset_display_name()
        )
        multimodal_dataset_operation = await self._create_multimodal_dataset(
            config=config,
            display_name=display_name,
            metadata_schema_uri=_datasets_utils.METADATA_SCHEMA_URI,
            metadata=multimodal_dataset.metadata,
        )
        response = await self._wait_for_operation(
            operation=multimodal_dataset_operation,
            timeout_seconds=config.timeout,
        )
        return _datasets_utils.create_from_response(
            types.MultimodalDataset, response, config
        )

    async def create_from_pandas(
        self,
        *,
        dataframe: pd.DataFrame,
        multimodal_dataset: Optional[types.MultimodalDatasetOrDict] = None,
        target_table_id: Optional[str] = None,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Creates a multimodal dataset from a pandas dataframe.

        Args:
            dataframe (pandas.DataFrame):
                The pandas dataframe to be used for the created dataset.
            multimodal_dataset:
                Optional. A representation of a multimodal dataset.
            target_table_id (str):
                Optional. The BigQuery table id where the dataframe will be
                uploaded. The table id can be in the format of "dataset.table"
                or "project.dataset.table". Note that the BigQuery
                dataset must already exist and be in the same location as the
                multimodal dataset. If not provided, a generated table id will
                be created in the `vertex_datasets` dataset (e.g.
                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
            config:
                Optional. A configuration for creating the multimodal dataset. If not
                provided, the default configuration will be used.

        Returns:
            dataset (MultimodalDataset):
                The created multimodal dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        elif not multimodal_dataset:
            multimodal_dataset = types.MultimodalDataset()

        bigframes = _datasets_utils._try_import_bigframes()
        project = self._api_client.project
        location = self._api_client.location
        credentials = self._api_client._credentials

        session_options = bigframes.BigQueryOptions(
            credentials=credentials,
            project=project,
            location=location,
        )
        with bigframes.connect(session_options) as session:
            return await self.create_from_bigframes(
                dataframe=session.read_pandas(dataframe),
                multimodal_dataset=multimodal_dataset,
                target_table_id=target_table_id,
                config=config,
            )

    async def create_from_bigframes(
        self,
        *,
        dataframe: "bigframes.pandas.DataFrame",  # type: ignore # noqa: F821
        multimodal_dataset: Optional[types.MultimodalDatasetOrDict] = None,
        target_table_id: Optional[str] = None,
        config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Creates a multimodal dataset from a bigframes dataframe.

        Args:
            dataframe (bigframes.pandas.DataFrame):
                The BigFrames dataframe that will be used for the created
                dataset.
            multimodal_dataset:
                Optional. A representation of a multimodal dataset.
            target_table_id (str):
                Optional. The BigQuery table id where the dataframe will be
                uploaded. The table id can be in the format of "dataset.table"
                or "project.dataset.table". Note that the BigQuery
                dataset must already exist and be in the same location as the
                multimodal dataset. If not provided, a generated table id will
                be created in the `vertex_datasets` dataset (e.g.
                `project.vertex_datasets_us_central1.multimodal_dataset_4cbf7ffd`).
            config:
                Optional. A configuration for creating the multimodal dataset. If not
                provided, the default configuration will be used.

        Returns:
            dataset (MultimodalDataset):
                The created multimodal dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        elif not multimodal_dataset:
            multimodal_dataset = types.MultimodalDataset()

        bigquery = _datasets_utils._try_import_bigquery()
        project = self._api_client.project
        location = self._api_client.location
        credentials = self._api_client._credentials

        if target_table_id:
            target_table_id = (
                await _datasets_utils._normalize_and_validate_table_id_async(
                    table_id=target_table_id,
                    project=project,
                    location=location,
                    credentials=credentials,
                )
            )
        else:
            dataset_id = await _datasets_utils._create_default_bigquery_dataset_if_not_exists_async(
                project=project, location=location, credentials=credentials
            )
            target_table_id = _datasets_utils._generate_target_table_id(dataset_id)

        client = bigquery.Client(project=project, credentials=credentials)
        await _datasets_utils.save_dataframe_to_bigquery_async(
            dataframe,
            target_table_id,
            client,
        )

        return await self.create_from_bigquery(
            multimodal_dataset=multimodal_dataset.model_copy(
                update={
                    "metadata": types.SchemaTablesDatasetMetadata(
                        input_config=types.SchemaTablesDatasetMetadataInputConfig(
                            bigquery_source=types.SchemaTablesDatasetMetadataBigQuerySource(
                                uri=f"bq://{target_table_id}"
                            )
                        )
                    )
                }
            ),
            config=config,
        )

    async def update_multimodal_dataset(
        self,
        *,
        multimodal_dataset: types.MultimodalDatasetOrDict,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Updates a multimodal dataset.

        Args:
          multimodal_dataset:
            Required. A representation of a multimodal dataset.
          config:
            Optional. A configuration for updating the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDataset object representing the updated multimodal
          dataset.
        """
        if isinstance(multimodal_dataset, dict):
            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
        _datasets_utils.validate_multimodal_dataset_bigquery_uri(multimodal_dataset)

        if isinstance(config, dict):
            config = types.VertexBaseConfig(**config)
        elif not config:
            config = types.VertexBaseConfig()

        return await self._update_multimodal_dataset(
            config=config,
            name=multimodal_dataset.name,
            display_name=multimodal_dataset.display_name,
            description=multimodal_dataset.description,
            metadata=multimodal_dataset.metadata,
        )

    async def get_multimodal_dataset(
        self,
        *,
        name: str,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDataset:
        """Gets a multimodal dataset.

        Args:
          name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          config:
            Optional. A configuration for getting the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDataset object representing the retrieved multimodal
          dataset.
        """
        if isinstance(config, dict):
            config = types.VertexBaseConfig(**config)
        elif not config:
            config = types.VertexBaseConfig()

        name = _datasets_utils.resolve_dataset_name(
            name, self._api_client.project, self._api_client.location
        )

        return await self._get_multimodal_dataset(config=config, name=name)

    async def delete_multimodal_dataset(
        self,
        *,
        name: str,
        config: Optional[types.VertexBaseConfigOrDict] = None,
    ) -> types.MultimodalDatasetOperation:
        """Deletes a multimodal dataset.

        Args:
          name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          config:
            Optional. A configuration for deleting the multimodal dataset. If not
            provided, the default configuration will be used.

        Returns:
          A types.MultimodalDatasetOperation object representing the delete
          multimodal dataset operation.
        """
        if isinstance(config, dict):
            config = types.VertexBaseConfig(**config)
        elif not config:
            config = types.VertexBaseConfig()

        name = _datasets_utils.resolve_dataset_name(
            name, self._api_client.project, self._api_client.location
        )

        return await self._delete_multimodal_dataset(config=config, name=name)

    async def assemble(
        self,
        *,
        name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssembleDatasetConfigOrDict] = None,
    ) -> str:
        """Assemble the dataset into a BigQuery table.

        Waits for the assemble operation to complete before returning.

        Args:
          name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          gemini_request_read_config:
            Optional. The read config to use to assemble the dataset. If
            not provided, the read config attached to the dataset will be
            used.
          config:
            Optional. A configuration for assembling the dataset. If not
            provided, the default configuration will be used.

        Returns:
            The URI of the bigquery table of the assembled dataset.
        """
        if isinstance(config, dict):
            config = types.AssembleDatasetConfig(**config)
        elif not config:
            config = types.AssembleDatasetConfig()

        name = _datasets_utils.resolve_dataset_name(
            name, self._api_client.project, self._api_client.location
        )

        operation = await self._assemble_multimodal_dataset(
            name=name,
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = await self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        return response["bigqueryDestination"]  # type: ignore[no-any-return]

    async def assess_tuning_resources(
        self,
        *,
        dataset_name: str,
        model_name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.TuningResourceUsageAssessmentResult:
        """Assess the tuning resources required for a given model.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
            Required. The name of the model to assess the tuning resources
            for.
          gemini_request_read_config:
            Optional. The read config used to assemble the dataset
            before assessing the tuning resources. If not provided, the
            read config attached to the dataset will be used. Required
            if no read config is attached to the dataset.
          config:
            Optional. A configuration for assessing the tuning resources. If not
            provided, the default configuration will be used.

        Returns:
          A types.TuningResourceUsageAssessmentResult object representing the
          tuning resource usage assessment result.
        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = await self._assess_multimodal_dataset(
            name=dataset_name,
            tuning_resource_usage_assessment_config=types.TuningResourceUsageAssessmentConfig(
                model_name=model_name
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = await self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        return _datasets_utils.create_from_response(
            types.TuningResourceUsageAssessmentResult,
            response["tuningResourceUsageAssessmentResult"],
            config,
        )

    async def assess_tuning_validity(
        self,
        *,
        dataset_name: str,
        model_name: str,
        dataset_usage: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.TuningValidationAssessmentResult:
        """Assess if the assembled dataset is valid in terms of tuning a given
        model.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
              Required. The name of the model to assess the tuning validity
              for.
          dataset_usage:
              Required. The dataset usage to assess the tuning validity for.
              Must be one of the following: SFT_TRAINING, SFT_VALIDATION.
          gemini_request_read_config:
              Optional. The read config used to assemble the dataset
              before assessing the tuning validity. If not provided, the
              read config attached to the dataset will be used. Required
              if no read config is attached to the dataset.
          config:
            Optional. A configuration for assessing the tuning validity. If not
            provided, the default configuration will be used.

        Returns:
            A dict containing the tuning validity assessment result. The dict
            contains the following keys:
            - errors: A list of errors that occurred during the tuning validity
              assessment.
        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = await self._assess_multimodal_dataset(
            name=dataset_name,
            tuning_validation_assessment_config=types.TuningValidationAssessmentConfig(
                model_name=model_name,
                dataset_usage=dataset_usage,
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = await self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        return _datasets_utils.create_from_response(
            types.TuningValidationAssessmentResult,
            response["tuningValidationAssessmentResult"],
            config,
        )

    async def assess_batch_prediction_resources(
        self,
        *,
        dataset_name: str,
        model_name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.BatchPredictionResourceUsageAssessmentResult:
        """Assess the batch prediction resources required for a given model.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
              Required. The name of the model to assess the batch prediction
              resources.
          gemini_request_read_config:
              Optional. The read config used to assemble the dataset
              before assessing the batch prediction resources. If not provided,
              the read config attached to the dataset will be used. Required
              if no read config is attached to the dataset.
          config:
              Optional. A configuration for assessing the batch prediction
              resources. If not provided, the default configuration will be
              used.

        Returns:
            A types.BatchPredictionResourceUsageAssessmentResult object
            representing the batch prediction resource usage assessment result.
            It contains the following keys:
              - token_count: The number of tokens in the dataset.
              - audio_token_count: The number of audio tokens in the dataset.

        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = await self._assess_multimodal_dataset(
            name=dataset_name,
            batch_prediction_resource_usage_assessment_config=types.BatchPredictionResourceUsageAssessmentConfig(
                model_name=model_name,
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = await self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        result = response["batchPredictionResourceUsageAssessmentResult"]
        return _datasets_utils.create_from_response(
            types.BatchPredictionResourceUsageAssessmentResult, result, config
        )

    async def assess_batch_prediction_validity(
        self,
        *,
        dataset_name: str,
        model_name: str,
        gemini_request_read_config: Optional[
            types.GeminiRequestReadConfigOrDict
        ] = None,
        config: Optional[types.AssessDatasetConfigOrDict] = None,
    ) -> types.BatchPredictionValidationAssessmentResult:
        """Assess if the assembled dataset is valid in terms of batch prediction
          for a given model. Raises an error if the dataset is invalid, otherwise
          returns None.

        Args:
          dataset_name:
            Required. A fully-qualified resource name or ID of the dataset.
            Example: "projects/.../locations/.../datasets/123" or "123".
          model_name:
            Required. The name of the model to assess the batch prediction
            validity for.
          gemini_request_read_config:
              Optional. The read config used to assemble the dataset
              before assessing the batch prediction validity. If not provided, the
              read config attached to the dataset will be used. Required
              if no read config is attached to the dataset.
          config:
            Optional. A configuration for assessing the batch prediction validity.
            If not provided, the default configuration will be used.

        Returns:
            A types.BatchPredictionValidationAssessmentResult object representing
            the batch prediction validity assessment result.
            It contains the following keys:
              - errors: A list of errors that occurred during the batch prediction
                validity assessment.
        """
        if isinstance(config, dict):
            config = types.AssessDatasetConfig(**config)
        elif not config:
            config = types.AssessDatasetConfig()

        dataset_name = _datasets_utils.resolve_dataset_name(
            dataset_name, self._api_client.project, self._api_client.location
        )

        operation = await self._assess_multimodal_dataset(
            name=dataset_name,
            batch_prediction_validation_assessment_config=types.BatchPredictionValidationAssessmentConfig(
                model_name=model_name,
            ),
            gemini_request_read_config=gemini_request_read_config,
            config=config,
        )
        response = await self._wait_for_operation(
            operation=operation,
            timeout_seconds=config.timeout,
        )
        result = response["batchPredictionValidationAssessmentResult"]
        return _datasets_utils.create_from_response(
            types.BatchPredictionValidationAssessmentResult, result, config
        )