From 34996a28de0495462aea9c7c2baf8b4249afb9df Mon Sep 17 00:00:00 2001 From: Ben Kazemi Date: Fri, 24 Oct 2025 12:08:23 -0700 Subject: [PATCH] feat: GenAI SDK client(multimodal) - Add get/update/list/delete to multimodal datasets. PiperOrigin-RevId: 823619345 --- .../test_create_multimodal_datasets.py | 78 ++- .../test_delete_multimodal_datasets.py | 71 +++ .../replays/test_get_multimodal_datasets.py | 50 ++ .../replays/test_list_multimodal_datasets.py | 41 ++ .../test_update_multimodal_datasets.py | 63 ++ vertexai/_genai/datasets.py | 573 +++++++++++++++++- vertexai/_genai/types/__init__.py | 26 + vertexai/_genai/types/common.py | 246 +++++++- 8 files changed, 1104 insertions(+), 44 deletions(-) create mode 100644 tests/unit/vertexai/genai/replays/test_delete_multimodal_datasets.py create mode 100644 tests/unit/vertexai/genai/replays/test_get_multimodal_datasets.py create mode 100644 tests/unit/vertexai/genai/replays/test_list_multimodal_datasets.py create mode 100644 tests/unit/vertexai/genai/replays/test_update_multimodal_datasets.py diff --git a/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py b/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py index 267dd5c629..d7608bc911 100644 --- a/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py +++ b/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py @@ -41,11 +41,32 @@ def test_create_dataset(client): def test_create_dataset_from_bigquery(client): - dataset = client.datasets.create_multimodal_dataset_from_bigquery( - multimodal_dataset=types.MultimodalDataset( - display_name="test-from-bigquery", - bigquery_uri=BIGQUERY_TABLE_NAME, - ) + dataset = client.datasets.create_from_bigquery( + multimodal_dataset={ + "display_name": "test-from-bigquery", + "description": "test-description-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + } + ) + assert isinstance(dataset, types.MultimodalDataset) + assert dataset.display_name == "test-from-bigquery" + + +def test_create_dataset_from_bigquery_without_bq_prefix(client): + dataset = client.datasets.create_from_bigquery( + multimodal_dataset={ + "display_name": "test-from-bigquery", + "description": "test-description-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": BIGQUERY_TABLE_NAME}, + }, + }, + }, ) assert isinstance(dataset, types.MultimodalDataset) assert dataset.display_name == "test-from-bigquery" @@ -77,11 +98,16 @@ async def test_create_dataset_async(client): @pytest.mark.asyncio async def test_create_dataset_from_bigquery_async(client): - dataset = await client.aio.datasets.create_multimodal_dataset_from_bigquery( - multimodal_dataset=types.MultimodalDataset( - display_name="test-from-bigquery", - bigquery_uri=BIGQUERY_TABLE_NAME, - ) + dataset = await client.aio.datasets.create_from_bigquery( + multimodal_dataset={ + "display_name": "test-from-bigquery", + "description": "test-description-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + } ) assert isinstance(dataset, types.MultimodalDataset) assert dataset.display_name == "test-from-bigquery" @@ -89,12 +115,34 @@ async def test_create_dataset_from_bigquery_async(client): @pytest.mark.asyncio async def test_create_dataset_from_bigquery_async_with_timeout(client): - dataset = await client.aio.datasets.create_multimodal_dataset_from_bigquery( + dataset = await client.aio.datasets.create_from_bigquery( config=types.CreateMultimodalDatasetConfig(timeout=120), - multimodal_dataset=types.MultimodalDataset( - display_name="test-from-bigquery", - bigquery_uri=BIGQUERY_TABLE_NAME, - ), + multimodal_dataset={ + "display_name": "test-from-bigquery", + "description": "test-description-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + }, + ) + assert isinstance(dataset, types.MultimodalDataset) + assert dataset.display_name == "test-from-bigquery" + + +@pytest.mark.asyncio +async def test_create_dataset_from_bigquery_async_without_bq_prefix(client): + dataset = await client.aio.datasets.create_from_bigquery( + multimodal_dataset={ + "display_name": "test-from-bigquery", + "description": "test-description-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": BIGQUERY_TABLE_NAME}, + }, + }, + }, ) assert isinstance(dataset, types.MultimodalDataset) assert dataset.display_name == "test-from-bigquery" diff --git a/tests/unit/vertexai/genai/replays/test_delete_multimodal_datasets.py b/tests/unit/vertexai/genai/replays/test_delete_multimodal_datasets.py new file mode 100644 index 0000000000..37216f94b7 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_delete_multimodal_datasets.py @@ -0,0 +1,71 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +import pytest + +BIGQUERY_TABLE_NAME = "vertex-sdk-dev.multimodal_dataset.test-table" + + +def test_delete_dataset(client): + dataset = client.datasets.create_from_bigquery( + multimodal_dataset={ + "display_name": "test-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + } + ) + name = dataset.name.split("/datasets/")[1] + + operation = client.datasets._delete_multimodal_dataset( + name=name, + ) + assert isinstance(operation, types.MultimodalDatasetOperation) + assert operation + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_delete_dataset_async(client): + dataset = client.datasets.create_from_bigquery( + multimodal_dataset={ + "display_name": "test-from-bigquery", + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + } + ) + name = dataset.name.split("/datasets/")[1] + + operation = client.datasets._delete_multimodal_dataset( + name=name, + ) + assert isinstance(operation, types.MultimodalDatasetOperation) + assert operation diff --git a/tests/unit/vertexai/genai/replays/test_get_multimodal_datasets.py b/tests/unit/vertexai/genai/replays/test_get_multimodal_datasets.py new file mode 100644 index 0000000000..8897a58039 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_get_multimodal_datasets.py @@ -0,0 +1,50 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +import pytest + +BIGQUERY_TABLE_NAME = "vertex-sdk-dev.multimodal_dataset.test-table" +DATASET = "8810841321427173376" + + +def test_get_dataset(client): + dataset = client.datasets._get_multimodal_dataset( + name=DATASET, + ) + assert isinstance(dataset, types.MultimodalDataset) + assert dataset.name.endswith(DATASET) + assert dataset.display_name == "test-from-bigquery" + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_get_dataset_async(client): + dataset = await client.aio.datasets._get_multimodal_dataset( + name=DATASET, + ) + assert isinstance(dataset, types.MultimodalDataset) + assert dataset.name.endswith(DATASET) + assert dataset.display_name == "test-from-bigquery" diff --git a/tests/unit/vertexai/genai/replays/test_list_multimodal_datasets.py b/tests/unit/vertexai/genai/replays/test_list_multimodal_datasets.py new file mode 100644 index 0000000000..1e71a53508 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_list_multimodal_datasets.py @@ -0,0 +1,41 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +import pytest + +BIGQUERY_TABLE_NAME = "vertex-sdk-dev.multimodal_dataset.test-table" + + +def test_list_dataset(client): + datasets = client.datasets._list_multimodal_datasets() + assert isinstance(datasets, types.ListMultimodalDatasetsResponse) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_list_dataset_async(client): + datasets = await client.aio.datasets._list_multimodal_datasets() + assert isinstance(datasets, types.ListMultimodalDatasetsResponse) diff --git a/tests/unit/vertexai/genai/replays/test_update_multimodal_datasets.py b/tests/unit/vertexai/genai/replays/test_update_multimodal_datasets.py new file mode 100644 index 0000000000..0af8d50f89 --- /dev/null +++ b/tests/unit/vertexai/genai/replays/test_update_multimodal_datasets.py @@ -0,0 +1,63 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from tests.unit.vertexai.genai.replays import pytest_helper +from vertexai._genai import types + +import pytest + +METADATA_SCHEMA_URI = ( + "gs://google-cloud-aiplatform/schema/dataset/metadata/multimodal_1.0.0.yaml" +) +BIGQUERY_TABLE_NAME = "vertex-sdk-dev.multimodal_dataset.test-table" +DATASET = "8810841321427173376" + + +def test_update_dataset(client): + operation = client.datasets._update_multimodal_dataset( + name=DATASET, + display_name="test-display-name", + description="test-description", + metadata={ + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + ) + assert isinstance(operation, types.MultimodalDatasetOperation) + + +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_update_dataset_async(client): + operation = await client.aio.datasets._update_multimodal_dataset( + name=DATASET, + display_name="test-display-name", + metadata={ + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + ) + assert isinstance(operation, types.MultimodalDatasetOperation) + assert operation diff --git a/vertexai/_genai/datasets.py b/vertexai/_genai/datasets.py index ecb39d6fc5..fb4ae6cc6e 100644 --- a/vertexai/_genai/datasets.py +++ b/vertexai/_genai/datasets.py @@ -66,6 +66,20 @@ def _CreateMultimodalDatasetParameters_to_vertex( 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, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + def _GetMultimodalDatasetOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -83,6 +97,81 @@ def _GetMultimodalDatasetOperationParameters_to_vertex( 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, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + 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, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + 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"])) + + return to_object + + class Datasets(_api_module.BaseModule): def _create_multimodal_dataset( @@ -148,6 +237,111 @@ def _create_multimodal_dataset( self._api_client._verify_response(return_value) return return_value + def _delete_multimodal_dataset( + self, *, config: Optional[types.VertexBaseConfigOrDict] = None, name: str + ) -> types.MultimodalDatasetOperation: + """ + Deletes a multimodal dataset resource. + """ + + parameter_model = types._DeleteMultimodalDatasetRequestParameters( + config=config, + name=name, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _DeleteMultimodalDatasetRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "datasets/{name}".format_map(request_url_dict) + else: + path = "datasets/{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("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=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get_multimodal_dataset( + self, + *, + config: Optional[types.VertexBaseConfigOrDict] = None, + name: Optional[str] = None, + ) -> types.MultimodalDataset: + """ + Gets a multimodal dataset resource. + """ + + parameter_model = types._GetMultimodalDatasetParameters( + config=config, + name=name, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _GetMultimodalDatasetParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "datasets/{name}".format_map(request_url_dict) + else: + path = "datasets/{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.MultimodalDataset._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + def _get_multimodal_dataset_operation( self, *, @@ -207,6 +401,118 @@ def _get_multimodal_dataset_operation( 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 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=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + + def _update_multimodal_dataset( + self, + *, + config: Optional[types.UpdateMultimodalDatasetConfigOrDict] = None, + name: Optional[str] = None, + display_name: Optional[str] = None, + metadata: Optional[dict[str, Any]] = None, + description: Optional[str] = None, + encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None, + ) -> types.MultimodalDatasetOperation: + """ + Updates a multimodal dataset resource. + """ + + parameter_model = types._UpdateMultimodalDatasetParameters( + config=config, + name=name, + display_name=display_name, + metadata=metadata, + description=description, + encryption_spec=encryption_spec, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _UpdateMultimodalDatasetParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "datasets/{name}".format_map(request_url_dict) + else: + path = "datasets/{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("patch", 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=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + def _wait_for_operation( self, operation: types.MultimodalDatasetOperation, @@ -271,10 +577,10 @@ def _wait_for_operation( ) return multimodal_operation.response - def create_multimodal_dataset_from_bigquery( + def create_from_bigquery( self, *, - multimodal_dataset: types.MultimodalDataset, + multimodal_dataset: types.MultimodalDatasetOrDict, config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None, ) -> types.MultimodalDataset: """Creates a multimodal dataset from a BigQuery table. @@ -289,8 +595,14 @@ def create_multimodal_dataset_from_bigquery( Returns: A types.MultimodalDataset object representing a multimodal dataset. """ - if not multimodal_dataset.bigquery_uri.startswith("bq://"): - multimodal_dataset.bigquery_uri = f"bq://{multimodal_dataset.bigquery_uri}" + if isinstance(multimodal_dataset, dict): + multimodal_dataset = types.MultimodalDataset(**multimodal_dataset) + if not multimodal_dataset.metadata.input_config.bigquery_source.uri.startswith( + "bq://" + ): + multimodal_dataset.metadata.input_config.bigquery_source.uri = ( + f"bq://{multimodal_dataset.metadata.input_config.bigquery_source.uri}" + ) if isinstance(config, dict): config = types.CreateMultimodalDatasetConfig(**config) elif not config: @@ -300,11 +612,7 @@ def create_multimodal_dataset_from_bigquery( config=config, display_name=multimodal_dataset.display_name, metadata_schema_uri=_datasets_utils.METADATA_SCHEMA_URI, - metadata={ - "inputConfig": { - "bigquerySource": {"uri": multimodal_dataset.bigquery_uri}, - }, - }, + metadata=multimodal_dataset.metadata, ) return self._wait_for_operation( operation=multimodal_dataset_operation, @@ -379,6 +687,115 @@ async def _create_multimodal_dataset( self._api_client._verify_response(return_value) return return_value + async def _delete_multimodal_dataset( + self, *, config: Optional[types.VertexBaseConfigOrDict] = None, name: str + ) -> types.MultimodalDatasetOperation: + """ + Deletes a multimodal dataset resource. + """ + + parameter_model = types._DeleteMultimodalDatasetRequestParameters( + config=config, + name=name, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _DeleteMultimodalDatasetRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "datasets/{name}".format_map(request_url_dict) + else: + path = "datasets/{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "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=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get_multimodal_dataset( + self, + *, + config: Optional[types.VertexBaseConfigOrDict] = None, + name: Optional[str] = None, + ) -> types.MultimodalDataset: + """ + Gets a multimodal dataset resource. + """ + + parameter_model = types._GetMultimodalDatasetParameters( + config=config, + name=name, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _GetMultimodalDatasetParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "datasets/{name}".format_map(request_url_dict) + else: + path = "datasets/{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.MultimodalDataset._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + async def _get_multimodal_dataset_operation( self, *, @@ -440,6 +857,122 @@ async def _get_multimodal_dataset_operation( 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 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=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update_multimodal_dataset( + self, + *, + config: Optional[types.UpdateMultimodalDatasetConfigOrDict] = None, + name: Optional[str] = None, + display_name: Optional[str] = None, + metadata: Optional[dict[str, Any]] = None, + description: Optional[str] = None, + encryption_spec: Optional[genai_types.EncryptionSpecOrDict] = None, + ) -> types.MultimodalDatasetOperation: + """ + Updates a multimodal dataset resource. + """ + + parameter_model = types._UpdateMultimodalDatasetParameters( + config=config, + name=name, + display_name=display_name, + metadata=metadata, + description=description, + encryption_spec=encryption_spec, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _UpdateMultimodalDatasetParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "datasets/{name}".format_map(request_url_dict) + else: + path = "datasets/{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", 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=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + async def _wait_for_operation( self, operation: types.MultimodalDatasetOperation, @@ -504,10 +1037,10 @@ async def _wait_for_operation( ) return multimodal_operation.response - async def create_multimodal_dataset_from_bigquery( + async def create_from_bigquery( self, *, - multimodal_dataset: types.MultimodalDataset, + multimodal_dataset: types.MultimodalDatasetOrDict, config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None, ) -> types.MultimodalDataset: """Creates a multimodal dataset from a BigQuery table. @@ -517,13 +1050,19 @@ async def create_multimodal_dataset_from_bigquery( Optional. A configuration for creating the multimodal dataset. If not provided, the default configuration will be used. multimodal_dataset: - Required. A representation of amultimodal dataset. + Required. A representation of a multimodal dataset. Returns: A types.MultimodalDataset object representing a multimodal dataset. """ - if not multimodal_dataset.bigquery_uri.startswith("bq://"): - multimodal_dataset.bigquery_uri = f"bq://{multimodal_dataset.bigquery_uri}" + if isinstance(multimodal_dataset, dict): + multimodal_dataset = types.MultimodalDataset(**multimodal_dataset) + if not multimodal_dataset.metadata.input_config.bigquery_source.uri.startswith( + "bq://" + ): + multimodal_dataset.metadata.input_config.bigquery_source.uri = ( + f"bq://{multimodal_dataset.metadata.input_config.bigquery_source.uri}" + ) if isinstance(config, dict): config = types.CreateMultimodalDatasetConfig(**config) elif not config: @@ -533,11 +1072,7 @@ async def create_multimodal_dataset_from_bigquery( config=config, display_name=multimodal_dataset.display_name, metadata_schema_uri=_datasets_utils.METADATA_SCHEMA_URI, - metadata={ - "inputConfig": { - "bigquerySource": {"uri": multimodal_dataset.bigquery_uri}, - }, - }, + metadata=multimodal_dataset.metadata, ) return await self._wait_for_operation( operation=multimodal_dataset_operation, diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index daa96925a6..714ad2634b 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -35,6 +35,7 @@ from .common import _DeleteAgentEngineSandboxRequestParameters from .common import _DeleteAgentEngineSessionRequestParameters from .common import _DeleteDatasetRequestParameters +from .common import _DeleteMultimodalDatasetRequestParameters from .common import _DeletePromptVersionRequestParameters from .common import _EvaluateInstancesRequestParameters from .common import _ExecuteCodeAgentEngineSandboxRequestParameters @@ -58,6 +59,7 @@ from .common import _GetEvaluationRunParameters from .common import _GetEvaluationSetParameters from .common import _GetMultimodalDatasetOperationParameters +from .common import _GetMultimodalDatasetParameters from .common import _ListAgentEngineMemoryRequestParameters from .common import _ListAgentEngineMemoryRevisionsRequestParameters from .common import _ListAgentEngineRequestParameters @@ -66,6 +68,7 @@ from .common import _ListAgentEngineSessionsRequestParameters from .common import _ListDatasetsRequestParameters from .common import _ListDatasetVersionsRequestParameters +from .common import _ListMultimodalDatasetsRequestParameters from .common import _OptimizeRequestParameters from .common import _QueryAgentEngineRequestParameters from .common import _RestoreVersionRequestParameters @@ -75,6 +78,7 @@ from .common import _UpdateAgentEngineRequestParameters from .common import _UpdateAgentEngineSessionRequestParameters from .common import _UpdateDatasetParameters +from .common import _UpdateMultimodalDatasetParameters from .common import AcceleratorType from .common import AgentEngine from .common import AgentEngineConfig @@ -461,6 +465,12 @@ from .common import ListDatasetVersionsResponse from .common import ListDatasetVersionsResponseDict from .common import ListDatasetVersionsResponseOrDict +from .common import ListMultimodalDatasetsConfig +from .common import ListMultimodalDatasetsConfigDict +from .common import ListMultimodalDatasetsConfigOrDict +from .common import ListMultimodalDatasetsResponse +from .common import ListMultimodalDatasetsResponseDict +from .common import ListMultimodalDatasetsResponseOrDict from .common import ListPromptsConfig from .common import ListPromptsConfigDict from .common import ListPromptsConfigOrDict @@ -935,6 +945,9 @@ from .common import UpdateDatasetConfig from .common import UpdateDatasetConfigDict from .common import UpdateDatasetConfigOrDict +from .common import UpdateMultimodalDatasetConfig +from .common import UpdateMultimodalDatasetConfigDict +from .common import UpdateMultimodalDatasetConfigOrDict from .common import VertexBaseConfig from .common import VertexBaseConfigDict from .common import VertexBaseConfigOrDict @@ -1591,6 +1604,15 @@ "GetMultimodalDatasetOperationConfig", "GetMultimodalDatasetOperationConfigDict", "GetMultimodalDatasetOperationConfigOrDict", + "ListMultimodalDatasetsConfig", + "ListMultimodalDatasetsConfigDict", + "ListMultimodalDatasetsConfigOrDict", + "ListMultimodalDatasetsResponse", + "ListMultimodalDatasetsResponseDict", + "ListMultimodalDatasetsResponseOrDict", + "UpdateMultimodalDatasetConfig", + "UpdateMultimodalDatasetConfigDict", + "UpdateMultimodalDatasetConfigOrDict", "CreateDatasetConfig", "CreateDatasetConfigDict", "CreateDatasetConfigOrDict", @@ -1838,7 +1860,11 @@ "_AppendAgentEngineSessionEventRequestParameters", "_ListAgentEngineSessionEventsRequestParameters", "_CreateMultimodalDatasetParameters", + "_DeleteMultimodalDatasetRequestParameters", + "_GetMultimodalDatasetParameters", "_GetMultimodalDatasetOperationParameters", + "_ListMultimodalDatasetsRequestParameters", + "_UpdateMultimodalDatasetParameters", "_CreateDatasetParameters", "_CreateDatasetVersionParameters", "_GetDatasetParameters", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index a27979b07d..5c386138ae 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -9544,7 +9544,7 @@ class SchemaTablesDatasetMetadataBigQuerySource(_common.BaseModel): uri: Optional[str] = Field( default=None, - description="""The URI of a BigQuery table. e.g. bq://projectId.bqDatasetId.bqTableId""", + description="""The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""", ) @@ -9552,7 +9552,7 @@ class SchemaTablesDatasetMetadataBigQuerySourceDict(TypedDict, total=False): """Represents the BigQuery source for multimodal dataset metadata.""" uri: Optional[str] - """The URI of a BigQuery table. e.g. bq://projectId.bqDatasetId.bqTableId""" + """The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""" SchemaTablesDatasetMetadataBigQuerySourceOrDict = Union[ @@ -9565,7 +9565,8 @@ class SchemaTablesDatasetMetadataInputConfig(_common.BaseModel): """Represents the input config for multimodal dataset metadata.""" bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySource] = Field( - default=None, description="""""" + default=None, + description="""The BigQuery source for multimodal dataset metadata.""", ) @@ -9573,7 +9574,7 @@ class SchemaTablesDatasetMetadataInputConfigDict(TypedDict, total=False): """Represents the input config for multimodal dataset metadata.""" bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySourceDict] - """""" + """The BigQuery source for multimodal dataset metadata.""" SchemaTablesDatasetMetadataInputConfigOrDict = Union[ @@ -9585,7 +9586,8 @@ class SchemaTablesDatasetMetadata(_common.BaseModel): """Represents the metadata schema for multimodal dataset metadata.""" input_config: Optional[SchemaTablesDatasetMetadataInputConfig] = Field( - default=None, description="""inputConfig""" + default=None, + description="""The input config for multimodal dataset metadata.""", ) @@ -9593,7 +9595,7 @@ class SchemaTablesDatasetMetadataDict(TypedDict, total=False): """Represents the metadata schema for multimodal dataset metadata.""" input_config: Optional[SchemaTablesDatasetMetadataInputConfigDict] - """inputConfig""" + """The input config for multimodal dataset metadata.""" SchemaTablesDatasetMetadataOrDict = Union[ @@ -9658,8 +9660,11 @@ class MultimodalDataset(_common.BaseModel): display_name: Optional[str] = Field( default=None, description="""The display name of the multimodal dataset.""" ) - bigquery_uri: Optional[str] = Field( - default=None, description="""The BigQuery URI of the multimodal dataset.""" + metadata: Optional[SchemaTablesDatasetMetadata] = Field( + default=None, description="""The metadata of the multimodal dataset.""" + ) + description: Optional[str] = Field( + default=None, description="""The description of the multimodal dataset.""" ) @@ -9672,8 +9677,11 @@ class MultimodalDatasetDict(TypedDict, total=False): display_name: Optional[str] """The display name of the multimodal dataset.""" - bigquery_uri: Optional[str] - """The BigQuery URI of the multimodal dataset.""" + metadata: Optional[SchemaTablesDatasetMetadataDict] + """The metadata of the multimodal dataset.""" + + description: Optional[str] + """The description of the multimodal dataset.""" MultimodalDatasetOrDict = Union[MultimodalDataset, MultimodalDatasetDict] @@ -9727,6 +9735,53 @@ class MultimodalDatasetOperationDict(TypedDict, total=False): ] +class _DeleteMultimodalDatasetRequestParameters(_common.BaseModel): + """Parameters for deleting a multimodal dataset.""" + + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + name: Optional[str] = Field( + default=None, description="""ID of the dataset to be deleted.""" + ) + + +class _DeleteMultimodalDatasetRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a multimodal dataset.""" + + config: Optional[VertexBaseConfigDict] + """""" + + name: Optional[str] + """ID of the dataset to be deleted.""" + + +_DeleteMultimodalDatasetRequestParametersOrDict = Union[ + _DeleteMultimodalDatasetRequestParameters, + _DeleteMultimodalDatasetRequestParametersDict, +] + + +class _GetMultimodalDatasetParameters(_common.BaseModel): + """Parameters for getting a multimodal dataset resource.""" + + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") + + +class _GetMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for getting a multimodal dataset resource.""" + + config: Optional[VertexBaseConfigDict] + """""" + + name: Optional[str] + """""" + + +_GetMultimodalDatasetParametersOrDict = Union[ + _GetMultimodalDatasetParameters, _GetMultimodalDatasetParametersDict +] + + class GetMultimodalDatasetOperationConfig(_common.BaseModel): """Config for getting a multimodal dataset operation.""" @@ -9776,6 +9831,177 @@ class _GetMultimodalDatasetOperationParametersDict(TypedDict, total=False): ] +class ListMultimodalDatasetsConfig(_common.BaseModel): + """Config for listing multimodal datasets.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( + default=None, + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", + ) + + +class ListMultimodalDatasetsConfigDict(TypedDict, total=False): + """Config for listing multimodal datasets.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" + + +ListMultimodalDatasetsConfigOrDict = Union[ + ListMultimodalDatasetsConfig, ListMultimodalDatasetsConfigDict +] + + +class _ListMultimodalDatasetsRequestParameters(_common.BaseModel): + """Parameters for listing multimodal datasets.""" + + config: Optional[ListMultimodalDatasetsConfig] = Field( + default=None, description="""""" + ) + + +class _ListMultimodalDatasetsRequestParametersDict(TypedDict, total=False): + """Parameters for listing multimodal datasets.""" + + config: Optional[ListMultimodalDatasetsConfigDict] + """""" + + +_ListMultimodalDatasetsRequestParametersOrDict = Union[ + _ListMultimodalDatasetsRequestParameters, + _ListMultimodalDatasetsRequestParametersDict, +] + + +class ListMultimodalDatasetsResponse(_common.BaseModel): + """Response for listing multimodal datasets.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the list datasets request in seconds. If not set, + the default timeout is 90 seconds.""", + ) + datasets: Optional[list[MultimodalDataset]] = Field( + default=None, + description="""List of datasets for the project. + """, + ) + + +class ListMultimodalDatasetsResponseDict(TypedDict, total=False): + """Response for listing multimodal datasets.""" + + sdk_http_response: Optional[genai_types.HttpResponseDict] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + timeout: Optional[int] + """The timeout for the list datasets request in seconds. If not set, + the default timeout is 90 seconds.""" + + datasets: Optional[list[MultimodalDatasetDict]] + """List of datasets for the project. + """ + + +ListMultimodalDatasetsResponseOrDict = Union[ + ListMultimodalDatasetsResponse, ListMultimodalDatasetsResponseDict +] + + +class UpdateMultimodalDatasetConfig(_common.BaseModel): + """Config for updating a multimodal dataset resource.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the update dataset request in seconds. If not set, + the default timeout is 90 seconds.""", + ) + + +class UpdateMultimodalDatasetConfigDict(TypedDict, total=False): + """Config for updating a multimodal dataset resource.""" + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + timeout: Optional[int] + """The timeout for the update dataset request in seconds. If not set, + the default timeout is 90 seconds.""" + + +UpdateMultimodalDatasetConfigOrDict = Union[ + UpdateMultimodalDatasetConfig, UpdateMultimodalDatasetConfigDict +] + + +class _UpdateMultimodalDatasetParameters(_common.BaseModel): + """Parameters for updating a multimodal dataset resource.""" + + config: Optional[UpdateMultimodalDatasetConfig] = Field( + default=None, description="""""" + ) + name: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + metadata: Optional[dict[str, Any]] = Field(default=None, description="""""") + description: Optional[str] = Field(default=None, description="""""") + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, description="""""" + ) + + +class _UpdateMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for updating a multimodal dataset resource.""" + + config: Optional[UpdateMultimodalDatasetConfigDict] + """""" + + name: Optional[str] + """""" + + display_name: Optional[str] + """""" + + metadata: Optional[dict[str, Any]] + """""" + + description: Optional[str] + """""" + + encryption_spec: Optional[genai_types.EncryptionSpecDict] + """""" + + +_UpdateMultimodalDatasetParametersOrDict = Union[ + _UpdateMultimodalDatasetParameters, _UpdateMultimodalDatasetParametersDict +] + + class CreateDatasetConfig(_common.BaseModel): """Config for creating a dataset resource to store prompts."""