diff --git a/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py b/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py index 78582bb8fa..9fa1711ac9 100644 --- a/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py +++ b/tests/unit/vertexai/genai/replays/test_create_multimodal_datasets.py @@ -72,6 +72,15 @@ def mock_import_bigframes(is_replay_mode): yield None +@pytest.fixture +def mock_generate_multimodal_dataset_display_name(): + with mock.patch.object( + _datasets_utils, "generate_multimodal_dataset_display_name" + ) as mock_generate: + mock_generate.return_value = "test-generated-name" + yield mock_generate + + def test_create_dataset(client): create_dataset_operation = client.datasets._create_multimodal_dataset( name="projects/vertex-sdk-dev/locations/us-central1", @@ -106,6 +115,21 @@ def test_create_dataset_from_bigquery(client): ) +@pytest.mark.usefixtures("mock_generate_multimodal_dataset_display_name") +def test_create_dataset_from_bigquery_no_display_name(client): + dataset = client.datasets.create_from_bigquery( + multimodal_dataset={ + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + } + ) + assert isinstance(dataset, types.MultimodalDataset) + assert dataset.display_name == "test-generated-name" + + @pytest.mark.usefixtures("mock_bigquery_client", "mock_import_bigframes") def test_create_dataset_from_pandas(client, is_replay_mode): dataframe = pd.DataFrame( @@ -230,6 +254,22 @@ async def test_create_dataset_from_bigquery_async(client): ) +@pytest.mark.asyncio +@pytest.mark.usefixtures("mock_generate_multimodal_dataset_display_name") +async def test_create_dataset_from_bigquery_no_display_name_async(client): + dataset = await client.aio.datasets.create_from_bigquery( + multimodal_dataset={ + "metadata": { + "inputConfig": { + "bigquerySource": {"uri": f"bq://{BIGQUERY_TABLE_NAME}"}, + }, + }, + } + ) + assert isinstance(dataset, types.MultimodalDataset) + assert dataset.display_name == "test-generated-name" + + @pytest.mark.asyncio async def test_create_dataset_from_bigquery_async_with_timeout(client): dataset = await client.aio.datasets.create_from_bigquery( diff --git a/vertexai/_genai/_datasets_utils.py b/vertexai/_genai/_datasets_utils.py index e49280f200..a913523e7e 100644 --- a/vertexai/_genai/_datasets_utils.py +++ b/vertexai/_genai/_datasets_utils.py @@ -15,6 +15,7 @@ """Utility functions for multimodal dataset.""" import asyncio +import datetime from typing import Any, Type, TypeVar import uuid @@ -227,6 +228,11 @@ def _generate_target_table_id(dataset_id: str) -> str: return f"{dataset_id}.{_DEFAULT_BQ_TABLE_PREFIX}_{str(uuid.uuid4())}" +def generate_multimodal_dataset_display_name() -> str: + """Generates a display name with a timestamp.""" + return f"MultimodalDataset {datetime.datetime.now().isoformat(sep=' ')}" + + def save_dataframe_to_bigquery( dataframe: "bigframes.pandas.DataFrame", # type: ignore # noqa: F821 target_table_id: str, diff --git a/vertexai/_genai/datasets.py b/vertexai/_genai/datasets.py index f5ef7793b1..0d9900e47e 100644 --- a/vertexai/_genai/datasets.py +++ b/vertexai/_genai/datasets.py @@ -933,9 +933,14 @@ def create_from_bigquery( elif not config: config = types.CreateMultimodalDatasetConfig() + display_name = ( + multimodal_dataset.display_name + if multimodal_dataset.display_name is not None + else _datasets_utils.generate_multimodal_dataset_display_name() + ) multimodal_dataset_operation = self._create_multimodal_dataset( config=config, - display_name=multimodal_dataset.display_name, + display_name=display_name, metadata_schema_uri=_datasets_utils.METADATA_SCHEMA_URI, metadata=multimodal_dataset.metadata, ) @@ -2150,9 +2155,14 @@ async def create_from_bigquery( elif not config: config = types.CreateMultimodalDatasetConfig() + display_name = ( + multimodal_dataset.display_name + if multimodal_dataset.display_name is not None + else _datasets_utils.generate_multimodal_dataset_display_name() + ) multimodal_dataset_operation = await self._create_multimodal_dataset( config=config, - display_name=multimodal_dataset.display_name, + display_name=display_name, metadata_schema_uri=_datasets_utils.METADATA_SCHEMA_URI, metadata=multimodal_dataset.metadata, )