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,37 @@ 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 +285,39 @@ 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_utils.py

@@ -34,17 +34,6 @@
 T = TypeVar("T", bound=BaseModel)
 
 
-def create_from_response(model_type: Type[T], response: dict[str, Any]) -> T:
-    """Creates a model from a response."""
-    model_field_names = model_type.model_fields.keys()
-    filtered_response = {}
-    for key, value in response.items():
-        snake_key = common.camel_to_snake(key)
-        if snake_key in model_field_names:
-            filtered_response[snake_key] = value
-    return model_type(**filtered_response)
-
-
 def validate_multimodal_dataset_bigquery_uri(
     multimodal_dataset: common.MultimodalDataset,
 ) -> None: