feat: GenAI SDK client(multimodal) - Add get/update/list/delete to multimodal datasets.

PiperOrigin-RevId: 823619345

Author: BenjaminKazemi

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,24 +98,51 @@ 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"
 
 
 @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"

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

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"

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)

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