feat: GenAI SDK client(multimodal) - Support getting a bigframe DataFrame from a multimodal dataset

PiperOrigin-RevId: 882583435

Author: cleop-google

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

@@ -123,6 +123,11 @@ def replays_prefix():
     return "test"
 
 
+@pytest.fixture
+def is_replay_mode(request):
+    return request.config.getoption("--mode") in ["replay", "tap"]
+
+
 @pytest.fixture
 def mock_agent_engine_create_path_exists():
     """Mocks os.path.exists to return True."""

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

@@ -28,11 +28,6 @@
 BIGQUERY_TABLE_NAME = "vertex-sdk-dev.multimodal_dataset.test-table"
 
 
-@pytest.fixture
-def is_replay_mode(request):
-    return request.config.getoption("--mode") in ["replay", "tap"]
-
-
 @pytest.fixture
 def mock_bigquery_client(is_replay_mode):
     if is_replay_mode:

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

@@ -15,14 +15,34 @@
 # 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,
@@ -41,6 +61,15 @@ 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(),
@@ -67,3 +96,13 @@ 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

vertexai/_genai/datasets.py

@@ -899,6 +899,33 @@ def create_from_pandas(
             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()
+
+        uri = _datasets_utils.multimodal_dataset_get_bigquery_uri(multimodal_dataset)
+        return bigframes.pandas.read_gbq_table(uri.removeprefix("bq://"))
+
     def update_multimodal_dataset(
         self,
         *,
@@ -1948,6 +1975,33 @@ async def create_from_pandas(
             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()
+
+        uri = _datasets_utils.multimodal_dataset_get_bigquery_uri(multimodal_dataset)
+        return bigframes.pandas.read_gbq_table(uri.removeprefix("bq://"))
+
     async def update_multimodal_dataset(
         self,
         *,