Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -254,6 +288,42 @@ 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):
Expand Down
13 changes: 1 addition & 12 deletions vertexai/_genai/_datasets_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

import asyncio
import datetime
from typing import Any, Type, TypeVar
from typing import Any, TypeVar
import uuid

import google.auth.credentials
Expand All @@ -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:
Expand Down
Loading
Loading