feat: GenAI SDK client(multimodal) - Accept an explicit bigquery_uri parameter in create_from_bigquery

PiperOrigin-RevId: 900174983

Author: cleop-google

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

@@ -115,6 +115,40 @@ def test_create_dataset_from_bigquery(client):
     )
 
 
+@pytest.mark.usefixtures("mock_generate_multimodal_dataset_display_name")
+def test_create_dataset_from_bigquery_with_uri(client):
+    dataset = client.datasets.create_from_bigquery(
+        bigquery_uri=f"bq://{BIGQUERY_TABLE_NAME}",
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.metadata.input_config.bigquery_source.uri == (
+        f"bq://{BIGQUERY_TABLE_NAME}"
+    )
+
+
+def test_create_dataset_from_bigquery_preserves_other_metadata(client):
+    dataset = client.datasets.create_from_bigquery(
+        bigquery_uri=f"bq://{BIGQUERY_TABLE_NAME}",
+        multimodal_dataset={
+            "display_name": "test-from-bigquery-uri",
+            "metadata": {
+                "gemini_request_read_config": {
+                    "assembled_request_column_name": "test_column"
+                }
+            },
+        },
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-from-bigquery-uri"
+    assert (
+        dataset.metadata.gemini_request_read_config.assembled_request_column_name
+        == "test_column"
+    )
+    assert dataset.metadata.input_config.bigquery_source.uri == (
+        f"bq://{BIGQUERY_TABLE_NAME}"
+    )
+
+
 @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(
@@ -254,6 +288,44 @@ 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_with_uri_async(client):
+    dataset = await client.aio.datasets.create_from_bigquery(
+        bigquery_uri=f"bq://{BIGQUERY_TABLE_NAME}",
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.metadata.input_config.bigquery_source.uri == (
+        f"bq://{BIGQUERY_TABLE_NAME}"
+    )
+
+
+@pytest.mark.asyncio
+async def test_create_dataset_from_bigquery_preserves_other_metadata_async(
+    client,
+):
+    dataset = await client.aio.datasets.create_from_bigquery(
+        bigquery_uri=f"bq://{BIGQUERY_TABLE_NAME}",
+        multimodal_dataset={
+            "display_name": "test-from-bigquery-uri",
+            "metadata": {
+                "gemini_request_read_config": {
+                    "assembled_request_column_name": "test_column"
+                }
+            },
+        },
+    )
+    assert isinstance(dataset, types.MultimodalDataset)
+    assert dataset.display_name == "test-from-bigquery-uri"
+    assert (
+        dataset.metadata.gemini_request_read_config.assembled_request_column_name
+        == "test_column"
+    )
+    assert dataset.metadata.input_config.bigquery_source.uri == (
+        f"bq://{BIGQUERY_TABLE_NAME}"
+    )
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("mock_generate_multimodal_dataset_display_name")
 async def test_create_dataset_from_bigquery_no_display_name_async(client):

vertexai/_genai/datasets.py

@@ -924,23 +924,34 @@ def _wait_for_operation(
     def create_from_bigquery(
         self,
         *,
-        multimodal_dataset: types.MultimodalDatasetOrDict,
+        bigquery_uri: Optional[str] = None,
+        multimodal_dataset: Optional[types.MultimodalDatasetOrDict] = None,
         config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
     ) -> types.MultimodalDataset:
         """Creates a multimodal dataset from a BigQuery table.
 
         Args:
+          bigquery_uri:
+            Optional. The BigQuery URI of the table to create the dataset from.
+            e.g. "bq://project.dataset.table".
           multimodal_dataset:
-            Required. A representation of a multimodal dataset.
+            Optional. A representation of a multimodal dataset.
           config:
             Optional. A configuration for creating the multimodal dataset. If not
             provided, the default configuration will be used.
 
         Returns:
           A types.MultimodalDataset object representing a multimodal dataset.
         """
-        if isinstance(multimodal_dataset, dict):
+        if multimodal_dataset is None:
+            multimodal_dataset = types.MultimodalDataset()
+        elif isinstance(multimodal_dataset, dict):
             multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
+
+        if bigquery_uri:
+            multimodal_dataset = multimodal_dataset.model_copy(deep=True)
+            multimodal_dataset.set_bigquery_uri(bigquery_uri)
+
         _datasets_utils.validate_multimodal_dataset_bigquery_uri(multimodal_dataset)
 
         if isinstance(config, dict):
@@ -2196,23 +2207,34 @@ async def _wait_for_operation(
     async def create_from_bigquery(
         self,
         *,
-        multimodal_dataset: types.MultimodalDatasetOrDict,
+        bigquery_uri: Optional[str] = None,
+        multimodal_dataset: Optional[types.MultimodalDatasetOrDict] = None,
         config: Optional[types.CreateMultimodalDatasetConfigOrDict] = None,
     ) -> types.MultimodalDataset:
         """Creates a multimodal dataset from a BigQuery table.
 
         Args:
+          bigquery_uri:
+            Optional. The BigQuery URI of the table to create the dataset from.
+            e.g. "bq://project.dataset.table".
           multimodal_dataset:
-            Required. A representation of a multimodal dataset.
+            Optional. A representation of a multimodal dataset.
           config:
             Optional. A configuration for creating the multimodal dataset. If not
             provided, the default configuration will be used.
 
         Returns:
           A types.MultimodalDataset object representing a multimodal dataset.
         """
-        if isinstance(multimodal_dataset, dict):
+        if multimodal_dataset is None:
+            multimodal_dataset = types.MultimodalDataset()
+        elif isinstance(multimodal_dataset, dict):
             multimodal_dataset = types.MultimodalDataset(**multimodal_dataset)
+
+        if bigquery_uri:
+            multimodal_dataset = multimodal_dataset.model_copy(deep=True)
+            multimodal_dataset.set_bigquery_uri(bigquery_uri)
+
         _datasets_utils.validate_multimodal_dataset_bigquery_uri(multimodal_dataset)
 
         if isinstance(config, dict):