chore: GenAI SDK client(multimodal) - Move to_bigframes method to MultimodalDataset class.

PiperOrigin-RevId: 889731146

Author: cleop-google

tests/unit/vertexai/genai/__init__.py

@@ -106,25 +106,6 @@ def test_create_dataset_from_bigquery(client):
     )
 
 
-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"
-    assert dataset.metadata.input_config.bigquery_source.uri == (
-        f"bq://{BIGQUERY_TABLE_NAME}"
-    )
-
-
 @pytest.mark.usefixtures("mock_bigquery_client", "mock_import_bigframes")
 def test_create_dataset_from_pandas(client, is_replay_mode):
     dataframe = pd.DataFrame(
@@ -270,26 +251,6 @@ async def test_create_dataset_from_bigquery_async_with_timeout(client):
     )
 
 
-@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"
-    assert dataset.metadata.input_config.bigquery_source.uri == (
-        f"bq://{BIGQUERY_TABLE_NAME}"
-    )
-
-
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("mock_bigquery_client", "mock_import_bigframes")
 async def test_create_dataset_from_pandas_async(client, is_replay_mode):

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

@@ -15,34 +15,14 @@
 # pylint: disable=protected-access,bad-continuation,missing-function-docstring
 
 from tests.unit.vertexai.genai.replays import pytest_helper
-from vertexai._genai import _datasets_utils
 from vertexai._genai import types
 
-from unittest import mock
 import pytest
 
 BIGQUERY_TABLE_NAME = "vertex-sdk-dev.multimodal_dataset.test-table"
 DATASET = "8810841321427173376"
 
 
-@pytest.fixture
-def mock_import_bigframes(is_replay_mode):
-    if is_replay_mode:
-        with mock.patch.object(
-            _datasets_utils, "_try_import_bigframes"
-        ) as mock_import_bigframes:
-            mock_read_gbq_table_result = mock.MagicMock()
-            mock_read_gbq_table_result.sql = f"SLECT * FROM `{BIGQUERY_TABLE_NAME}`"
-
-            bigframes = mock.MagicMock()
-            bigframes.pandas.read_gbq_table.return_value = mock_read_gbq_table_result
-
-            mock_import_bigframes.return_value = bigframes
-            yield mock_import_bigframes
-    else:
-        yield None
-
-
 def test_get_dataset(client):
     dataset = client.datasets._get_multimodal_dataset(
         name=DATASET,
@@ -61,15 +41,6 @@ def test_get_dataset_from_public_method(client):
     assert dataset.display_name == "test-display-name"
 
 
-@pytest.mark.usefixtures("mock_import_bigframes")
-def test_to_bigframes(client):
-    dataset = client.datasets.get_multimodal_dataset(
-        name=DATASET,
-    )
-    df = client.datasets.to_bigframes(multimodal_dataset=dataset)
-    assert BIGQUERY_TABLE_NAME in df.sql
-
-
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),
@@ -96,13 +67,3 @@ async def test_get_dataset_from_public_method_async(client):
     assert isinstance(dataset, types.MultimodalDataset)
     assert dataset.name.endswith(DATASET)
     assert dataset.display_name == "test-display-name"
-
-
-@pytest.mark.asyncio
-@pytest.mark.usefixtures("mock_import_bigframes")
-async def test_to_bigframes_async(client):
-    dataset = await client.aio.datasets.get_multimodal_dataset(
-        name=DATASET,
-    )
-    df = await client.aio.datasets.to_bigframes(multimodal_dataset=dataset)
-    assert BIGQUERY_TABLE_NAME in df.sql

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

@@ -0,0 +1,159 @@
+# 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.
+#
+"""Tests for multimodal datasets."""
+from unittest import mock
+
+from vertexai._genai import _datasets_utils
+from vertexai._genai import types
+import pytest
+
+
+@pytest.fixture
+def mock_import_bigframes():
+    with mock.patch.object(
+        _datasets_utils, "_try_import_bigframes"
+    ) as mock_import_bigframes:
+        mock_read_gbq_table_result = mock.MagicMock()
+        mock_read_gbq_table_result.sql = "SELECT * FROM `project.dataset.table`"
+
+        bigframes = mock.MagicMock()
+        bigframes.pandas.read_gbq_table.return_value = mock_read_gbq_table_result
+
+        mock_import_bigframes.return_value = bigframes
+        yield mock_import_bigframes
+
+
+class TestMultimodalDataset:
+
+    def test_get_read_config(self):
+        dataset = types.MultimodalDataset(
+            metadata={
+                "gemini_request_read_config": {
+                    "assembled_request_column_name": "test_column",
+                },
+            },
+        )
+
+        read_config = dataset.get_read_config()
+
+        assert isinstance(read_config, types.GeminiRequestReadConfig)
+        assert read_config.assembled_request_column_name == "test_column"
+
+    def test_get_read_config_empty(self):
+        dataset = types.MultimodalDataset()
+        assert dataset.get_read_config() is None
+
+    def test_set_read_config(self):
+        dataset = types.MultimodalDataset()
+
+        dataset.set_read_config(
+            read_config={
+                "assembled_request_column_name": "test_column",
+            },
+        )
+
+        assert isinstance(dataset, types.MultimodalDataset)
+        assert (
+            dataset.metadata.gemini_request_read_config.assembled_request_column_name
+            == "test_column"
+        )
+
+    def test_set_read_config_preserves_other_fields(self):
+        dataset = types.MultimodalDataset(
+            metadata={
+                "inputConfig": {
+                    "bigquerySource": {"uri": "bq://test_table"},
+                },
+            },
+        )
+
+        dataset.set_read_config(
+            read_config={
+                "assembled_request_column_name": "test_column",
+            },
+        )
+
+        assert isinstance(dataset, types.MultimodalDataset)
+        assert (
+            dataset.metadata.gemini_request_read_config.assembled_request_column_name
+            == "test_column"
+        )
+        assert dataset.metadata.input_config.bigquery_source.uri == "bq://test_table"
+
+    def test_get_bigquery_uri(self):
+        dataset = types.MultimodalDataset(
+            metadata={
+                "inputConfig": {
+                    "bigquerySource": {"uri": "bq://project.dataset.table"},
+                },
+            },
+        )
+
+        assert dataset.get_bigquery_uri() == "bq://project.dataset.table"
+
+    def test_get_bigquery_uri_empty(self):
+        dataset = types.MultimodalDataset()
+        assert dataset.get_bigquery_uri() is None
+
+    def test_set_bigquery_uri(self):
+        dataset = types.MultimodalDataset()
+
+        dataset.set_bigquery_uri("bq://project.dataset.table")
+
+        assert isinstance(dataset, types.MultimodalDataset)
+        assert (
+            dataset.metadata.input_config.bigquery_source.uri
+            == "bq://project.dataset.table"
+        )
+
+    def test_set_bigquery_uri_without_prefix(self):
+        dataset = types.MultimodalDataset()
+
+        dataset.set_bigquery_uri("project.dataset.table")
+
+        assert isinstance(dataset, types.MultimodalDataset)
+        assert (
+            dataset.metadata.input_config.bigquery_source.uri
+            == "bq://project.dataset.table"
+        )
+
+    def test_set_bigquery_uri_preserves_other_fields(self):
+        dataset = types.MultimodalDataset(
+            metadata={
+                "gemini_request_read_config": {
+                    "assembled_request_column_name": "test_column",
+                },
+            },
+        )
+
+        dataset.set_bigquery_uri("bq://test_table")
+
+        assert isinstance(dataset, types.MultimodalDataset)
+        assert dataset.metadata.input_config.bigquery_source.uri == "bq://test_table"
+        assert (
+            dataset.metadata.gemini_request_read_config.assembled_request_column_name
+            == "test_column"
+        )
+
+    def test_to_bigframes(self, mock_import_bigframes):
+        dataset = types.MultimodalDataset()
+        dataset.set_bigquery_uri("bq://project.dataset.table")
+
+        df = dataset.to_bigframes()
+
+        assert "project.dataset.table" in df.sql
+        mock_import_bigframes.return_value.pandas.read_gbq_table.assert_called_once_with(
+            "project.dataset.table"
+        )

tests/unit/vertexai/genai/test_datasets.py

@@ -44,10 +44,10 @@ def create_from_response(model_type: Type[T], response: dict[str, Any]) -> T:
     return model_type(**filtered_response)
 
 
-def multimodal_dataset_get_bigquery_uri(
+def validate_multimodal_dataset_bigquery_uri(
     multimodal_dataset: common.MultimodalDataset,
-) -> str:
-    """Gets the bigquery uri from a multimodal dataset or raises ValueError."""
+) -> None:
+    """Validates that a multimodal dataset has a bigquery uri or raises ValueError."""
     if (
         not hasattr(multimodal_dataset, "metadata")
         or multimodal_dataset.metadata is None
@@ -70,33 +70,12 @@ def multimodal_dataset_get_bigquery_uri(
         raise ValueError(
             "Multimodal dataset input config bigquery source uri is required."
         )
-    return str(multimodal_dataset.metadata.input_config.bigquery_source.uri)
-
-
-def multimodal_dataset_set_bigquery_uri(
-    multimodal_dataset: common.MultimodalDataset,
-    bigquery_uri: str,
-) -> None:
-    """Sets the bigquery uri from a multimodal dataset or raises ValueError."""
-    metadata = (
-        common.SchemaTablesDatasetMetadata()
-        if multimodal_dataset.metadata is None
-        else multimodal_dataset.metadata
-    )
-    input_config = (
-        common.SchemaTablesDatasetMetadataInputConfig()
-        if metadata.input_config is None
-        else metadata.input_config
-    )
-    bigquery_source = (
-        common.SchemaTablesDatasetMetadataBigQuerySource()
-        if input_config.bigquery_source is None
-        else input_config.bigquery_source
-    )
-    bigquery_source.uri = bigquery_uri
-    input_config.bigquery_source = bigquery_source
-    metadata.input_config = input_config
-    multimodal_dataset.metadata = metadata
+    if not str(multimodal_dataset.metadata.input_config.bigquery_source.uri).startswith(
+        "bq://"
+    ):
+        raise ValueError(
+            "Multimodal dataset bigquery source uri must start with 'bq://'."
+        )
 
 
 def _try_import_bigframes() -> Any: