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

BREAKING CHANGE: `to_bigframes` has been removed from the datasets module and moved into the `MultimodalDataset` class.

PiperOrigin-RevId: 889731146

Author: cleop-google

tests/unit/vertexai/genai/replays/test_get_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/test_multimodal_datasets_genai.py

@@ -13,8 +13,26 @@
 # 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:
@@ -126,3 +144,14 @@ def test_set_bigquery_uri_preserves_other_fields(self):
             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"
+        )

vertexai/_genai/datasets.py

@@ -970,36 +970,6 @@ def create_from_bigframes(
             config=config,
         )
 
-    def to_bigframes(
-        self,
-        *,
-        multimodal_dataset: types.MultimodalDatasetOrDict,
-    ) -> "bigframes.pandas.DataFrame":  # type: ignore # noqa: F821
-        """Converts a multimodal dataset to a BigFrames dataframe.
-
-        This is the preferred method to inspect the multimodal dataset in a
-        notebook.
-
-        Args:
-          multimodal_dataset:
-            Required. A representation of a multimodal dataset.
-
-        Returns:
-          A BigFrames dataframe.
-        """
-        bigframes = _datasets_utils._try_import_bigframes()
-
-        if isinstance(multimodal_dataset, dict):
-            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
-        elif not multimodal_dataset:
-            multimodal_dataset = types.MultimodalDataset()
-
-        if multimodal_dataset.bigquery_uri is None:
-            raise ValueError("Multimodal dataset bigquery source uri is not set.")
-        return bigframes.pandas.read_gbq_table(
-            multimodal_dataset.bigquery_uri.removeprefix("bq://")
-        )
-
     def update_multimodal_dataset(
         self,
         *,
@@ -2115,37 +2085,6 @@ async def create_from_bigframes(
             config=config,
         )
 
-    async def to_bigframes(
-        self,
-        *,
-        multimodal_dataset: types.MultimodalDatasetOrDict,
-    ) -> "bigframes.pandas.DataFrame":  # type: ignore # noqa: F821
-        """Converts a multimodal dataset to a BigFrames dataframe.
-
-        This is the preferred method to inspect the multimodal dataset in a
-        notebook.
-
-        Args:
-          multimodal_dataset:
-            Required. A representation of a multimodal dataset.
-
-        Returns:
-          A BigFrames dataframe.
-        """
-        bigframes = _datasets_utils._try_import_bigframes()
-
-        if isinstance(multimodal_dataset, dict):
-            multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
-        elif not multimodal_dataset:
-            multimodal_dataset = types.MultimodalDataset()
-
-        if multimodal_dataset.bigquery_uri is None:
-            raise ValueError("Multimodal dataset bigquery source uri is missing.")
-        return await asyncio.to_thread(
-            bigframes.pandas.read_gbq_table,
-            multimodal_dataset.bigquery_uri.removeprefix("bq://"),
-        )
-
     async def update_multimodal_dataset(
         self,
         *,

vertexai/_genai/types/common.py

@@ -12424,6 +12424,25 @@ def set_bigquery_uri(
         metadata.input_config = input_config
         self.metadata = metadata
 
+    def to_bigframes(
+        self,
+    ) -> "bigframes.pandas.DataFrame":  # type: ignore # noqa: F821
+        """Converts the multimodal dataset to a BigFrames dataframe.
+
+        This is the preferred method to inspect the multimodal dataset in a
+        notebook.
+
+        Returns:
+          A BigFrames dataframe.
+        """
+        from .. import _datasets_utils
+
+        bigframes = _datasets_utils._try_import_bigframes()
+
+        if self.bigquery_uri is None:
+            raise ValueError("Multimodal dataset bigquery source uri is not set.")
+        return bigframes.pandas.read_gbq_table(self.bigquery_uri.removeprefix("bq://"))
+
 
 class MultimodalDatasetDict(TypedDict, total=False):
     """Represents a multimodal dataset."""