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

PiperOrigin-RevId: 825617246

Author: BenjaminKazemi

tests/unit/vertexai/genai/replays/test_delete_multimodal_datasets.py

@@ -39,7 +39,27 @@ def test_delete_dataset(client):
         name=name,
     )
     assert isinstance(operation, types.MultimodalDatasetOperation)
-    assert operation
+    assert operation.done
+
+
+def test_delete_dataset_with_public_method(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.done
 
 
 pytestmark = pytest_helper.setup(
@@ -52,7 +72,7 @@ def test_delete_dataset(client):
 
 @pytest.mark.asyncio
 async def test_delete_dataset_async(client):
-    dataset = client.datasets.create_from_bigquery(
+    dataset = await client.aio.datasets.create_from_bigquery(
         multimodal_dataset={
             "display_name": "test-from-bigquery",
             "metadata": {
@@ -64,8 +84,29 @@ async def test_delete_dataset_async(client):
     )
     name = dataset.name.split("/datasets/")[1]
 
-    operation = client.datasets._delete_multimodal_dataset(
+    operation = await client.aio.datasets._delete_multimodal_dataset(
         name=name,
     )
     assert isinstance(operation, types.MultimodalDatasetOperation)
     assert operation
+
+
+@pytest.mark.asyncio
+async def test_delete_dataset_with_public_method_async(client):
+    dataset = await client.aio.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 = await client.aio.datasets.delete_multimodal_dataset(
+        name=name,
+    )
+    assert isinstance(operation, types.MultimodalDatasetOperation)
+    assert operation.done

tests/unit/vertexai/genai/replays/test_get_multimodal_datasets.py

@@ -29,7 +29,16 @@ def test_get_dataset(client):
     )
     assert isinstance(dataset, types.MultimodalDataset)
     assert dataset.name.endswith(DATASET)
-    assert dataset.display_name == "test-from-bigquery"
+    assert dataset.display_name == "test-display-name"
+
+
+def test_get_dataset_from_public_method(client):
+    dataset = client.datasets.get_multimodal_dataset(
+        name=DATASET,
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.name.endswith(DATASET)
+    assert dataset.display_name == "test-display-name"
 
 
 pytestmark = pytest_helper.setup(
@@ -47,4 +56,14 @@ async def test_get_dataset_async(client):
     )
     assert isinstance(dataset, types.MultimodalDataset)
     assert dataset.name.endswith(DATASET)
-    assert dataset.display_name == "test-from-bigquery"
+    assert dataset.display_name == "test-display-name"
+
+
+@pytest.mark.asyncio
+async def test_get_dataset_from_public_method_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-display-name"

tests/unit/vertexai/genai/replays/test_update_multimodal_datasets.py

@@ -27,17 +27,37 @@
 
 
 def test_update_dataset(client):
-    operation = client.datasets._update_multimodal_dataset(
+    dataset = client.datasets._update_multimodal_dataset(
         name=DATASET,
-        display_name="test-display-name",
-        description="test-description",
+        display_name="test-display-name (updated with internal method)",
+        description="test-description (updated with internal method)",
         metadata={
             "inputConfig": {
                 "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"},
             },
         },
     )
-    assert isinstance(operation, types.MultimodalDatasetOperation)
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-display-name (updated with internal method)"
+    assert dataset.description == "test-description (updated with internal method)"
+
+
+def test_update_dataset_with_public_method(client):
+    dataset = client.datasets.update_multimodal_dataset(
+        multimodal_dataset={
+            "name": DATASET,
+            "display_name": "test-display-name (updated with public method)",
+            "description": "test-description (updated with public method)",
+            "metadata": {
+                "inputConfig": {
+                    "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"},
+                },
+            },
+        }
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-display-name (updated with public method)"
+    assert dataset.description == "test-description (updated with public method)"
 
 
 pytestmark = pytest_helper.setup(
@@ -50,14 +70,35 @@ def test_update_dataset(client):
 
 @pytest.mark.asyncio
 async def test_update_dataset_async(client):
-    operation = await client.aio.datasets._update_multimodal_dataset(
+    dataset = await client.aio.datasets._update_multimodal_dataset(
         name=DATASET,
-        display_name="test-display-name",
+        display_name="test-display-name (updated with internal method)",
+        description="test-description (updated with internal method)",
         metadata={
             "inputConfig": {
                 "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"},
             },
         },
     )
-    assert isinstance(operation, types.MultimodalDatasetOperation)
-    assert operation
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-display-name (updated with internal method)"
+    assert dataset.description == "test-description (updated with internal method)"
+
+
+@pytest.mark.asyncio
+async def test_update_dataset_with_public_method_async(client):
+    dataset = await client.aio.datasets.update_multimodal_dataset(
+        multimodal_dataset={
+            "name": DATASET,
+            "display_name": "test-display-name (updated with public method)",
+            "description": "test-description (updated with public method)",
+            "metadata": {
+                "inputConfig": {
+                    "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"},
+                },
+            },
+        }
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-display-name (updated with public method)"
+    assert dataset.description == "test-description (updated with public method)"