Skip to content

Commit fafcaf0

Browse files
cleop-googlecopybara-github
cleop-google
authored and
copybara-github
committed
feat: GenAI SDK client(multimodal) - Add to_batch_job_source and get_batch_job_destination to MultimodalDataset
PiperOrigin-RevId: 906352851
1 parent 2faf725 commit fafcaf0

4 files changed

Lines changed: 134 additions & 56 deletions

File tree

‎agentplatform/_genai/_datasets_utils.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,13 @@ def generate_multimodal_dataset_display_name() -> str:
242242
return f"MultimodalDataset {datetime.datetime.now().isoformat(sep=' ')}"
243243

244244

245+
def get_batch_job_unique_name() -> str:
246+
"""Generates a unique name suffix for a batch job destination."""
247+
timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
248+
unique_id = uuid.uuid4().hex[0:5]
249+
return f"{timestamp}_{unique_id}"
250+
251+
245252
def save_dataframe_to_bigquery(
246253
dataframe: "bigframes.pandas.DataFrame", # type: ignore # noqa: F821
247254
target_table_id: str,

‎agentplatform/_genai/types/common.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18722,6 +18722,28 @@ def to_bigframes(
1872218722
raise ValueError("Multimodal dataset bigquery source uri is not set.")
1872318723
return bigframes.pandas.read_gbq_table(self.bigquery_uri.removeprefix("bq://"))
1872418724

18725+
def to_batch_job_source(self) -> "genai_types.BatchJobSource":
18726+
"""Converts the dataset to a BatchJobSource."""
18727+
return genai_types.BatchJobSource(
18728+
vertex_dataset_name=self.name,
18729+
)
18730+
18731+
def get_batch_job_destination(self) -> "genai_types.BatchJobDestination":
18732+
"""Converts the dataset to a BatchJobDestination."""
18733+
from .. import _datasets_utils
18734+
18735+
unique_name = _datasets_utils.get_batch_job_unique_name()
18736+
bigquery_uri = self.bigquery_uri
18737+
if bigquery_uri is None:
18738+
raise ValueError("Multimodal dataset bigquery source uri is not set.")
18739+
curr_display_name = self.display_name or "genai_batch_job"
18740+
return genai_types.BatchJobDestination(
18741+
vertex_dataset=genai_types.VertexMultimodalDatasetDestination(
18742+
display_name=f"{curr_display_name}_batch_output_{unique_name}",
18743+
bigquery_destination=f"{bigquery_uri}_batch_output_{unique_name}",
18744+
)
18745+
)
18746+
1872518747

1872618748
class MultimodalDatasetDict(TypedDict, total=False):
1872718749
"""Represents a multimodal dataset."""

‎tests/unit/agentplatform/genai/replays/test_create_multimodal_datasets.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,11 +74,20 @@ def mock_import_bigframes(is_replay_mode):
7474

7575
@pytest.fixture
7676
def mock_generate_multimodal_dataset_display_name():
77-
with mock.patch.object(
77+
with mock.patch.object(
7878
_datasets_utils, "generate_multimodal_dataset_display_name"
7979
) as mock_generate:
80-
mock_generate.return_value = "test-generated-name"
81-
yield mock_generate
80+
mock_generate.return_value = "test-generated-name"
81+
yield mock_generate
82+
83+
84+
@pytest.fixture
85+
def mock_get_batch_job_unique_name():
86+
with mock.patch.object(
87+
_datasets_utils, "get_batch_job_unique_name"
88+
) as mock_unique_name:
89+
mock_unique_name.return_value = "12345678901234_abcde"
90+
yield mock_unique_name
8291

8392

8493
def test_create_dataset(client):

‎tests/unit/agentplatform/genai/test_multimodal_datasets_genai.py‎

Lines changed: 93 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -23,140 +23,180 @@
2323

2424
@pytest.fixture
2525
def mock_import_bigframes():
26-
with mock.patch.object(
26+
with mock.patch.object(
2727
_datasets_utils, "_try_import_bigframes"
2828
) as mock_import_bigframes:
29-
mock_read_gbq_table_result = mock.MagicMock()
30-
mock_read_gbq_table_result.sql = "SELECT * FROM `project.dataset.table`"
29+
mock_read_gbq_table_result = mock.MagicMock()
30+
mock_read_gbq_table_result.sql = "SELECT * FROM `project.dataset.table`"
3131

32-
bigframes = mock.MagicMock()
33-
bigframes.pandas.read_gbq_table.return_value = mock_read_gbq_table_result
32+
bigframes = mock.MagicMock()
33+
bigframes.pandas.read_gbq_table.return_value = mock_read_gbq_table_result
3434

35-
mock_import_bigframes.return_value = bigframes
36-
yield mock_import_bigframes
35+
mock_import_bigframes.return_value = bigframes
36+
yield mock_import_bigframes
37+
38+
39+
@pytest.fixture
40+
def mock_get_batch_job_unique_name():
41+
with mock.patch.object(
42+
_datasets_utils, "get_batch_job_unique_name"
43+
) as mock_unique_name:
44+
mock_unique_name.return_value = "12345678901234_abcde"
45+
yield mock_unique_name
3746

3847

3948
class TestMultimodalDataset:
4049

41-
def test_read_config(self):
42-
dataset = types.MultimodalDataset(
50+
def test_read_config(self):
51+
dataset = types.MultimodalDataset(
4352
metadata={
4453
"gemini_request_read_config": {
4554
"assembled_request_column_name": "test_column",
4655
},
4756
},
4857
)
4958

50-
assert isinstance(dataset.read_config, types.GeminiRequestReadConfig)
51-
assert dataset.read_config.assembled_request_column_name == "test_column"
59+
assert isinstance(dataset.read_config, types.GeminiRequestReadConfig)
60+
assert dataset.read_config.assembled_request_column_name == "test_column"
5261

53-
def test_read_config_empty(self):
54-
dataset = types.MultimodalDataset()
55-
assert dataset.read_config is None
62+
def test_read_config_empty(self):
63+
dataset = types.MultimodalDataset()
64+
assert dataset.read_config is None
5665

57-
def test_set_read_config(self):
58-
dataset = types.MultimodalDataset()
66+
def test_set_read_config(self):
67+
dataset = types.MultimodalDataset()
5968

60-
dataset.set_read_config(
69+
dataset.set_read_config(
6170
read_config={
6271
"assembled_request_column_name": "test_column",
6372
},
6473
)
6574

66-
assert isinstance(dataset, types.MultimodalDataset)
67-
assert (
75+
assert isinstance(dataset, types.MultimodalDataset)
76+
assert (
6877
dataset.metadata.gemini_request_read_config.assembled_request_column_name
6978
== "test_column"
7079
)
7180

72-
def test_set_read_config_preserves_other_fields(self):
73-
dataset = types.MultimodalDataset(
81+
def test_set_read_config_preserves_other_fields(self):
82+
dataset = types.MultimodalDataset(
7483
metadata={
7584
"inputConfig": {
7685
"bigquerySource": {"uri": "bq://test_table"},
7786
},
7887
},
7988
)
8089

81-
dataset.set_read_config(
90+
dataset.set_read_config(
8291
read_config={
8392
"assembled_request_column_name": "test_column",
8493
},
8594
)
8695

87-
assert isinstance(dataset, types.MultimodalDataset)
88-
assert (
96+
assert isinstance(dataset, types.MultimodalDataset)
97+
assert (
8998
dataset.metadata.gemini_request_read_config.assembled_request_column_name
9099
== "test_column"
91100
)
92-
assert dataset.metadata.input_config.bigquery_source.uri == "bq://test_table"
101+
assert dataset.metadata.input_config.bigquery_source.uri == "bq://test_table"
93102

94-
def test_bigquery_uri(self):
95-
dataset = types.MultimodalDataset(
103+
def test_bigquery_uri(self):
104+
dataset = types.MultimodalDataset(
96105
metadata={
97106
"inputConfig": {
98107
"bigquerySource": {"uri": "bq://project.dataset.table"},
99108
},
100109
},
101110
)
102111

103-
assert dataset.bigquery_uri == "bq://project.dataset.table"
112+
assert dataset.bigquery_uri == "bq://project.dataset.table"
104113

105-
def test_bigquery_uri_empty(self):
106-
dataset = types.MultimodalDataset()
107-
assert dataset.bigquery_uri is None
114+
def test_bigquery_uri_empty(self):
115+
dataset = types.MultimodalDataset()
116+
assert dataset.bigquery_uri is None
108117

109-
def test_set_bigquery_uri(self):
110-
dataset = types.MultimodalDataset()
118+
def test_set_bigquery_uri(self):
119+
dataset = types.MultimodalDataset()
111120

112-
dataset.set_bigquery_uri("bq://project.dataset.table")
121+
dataset.set_bigquery_uri("bq://project.dataset.table")
113122

114-
assert isinstance(dataset, types.MultimodalDataset)
115-
assert (
123+
assert isinstance(dataset, types.MultimodalDataset)
124+
assert (
116125
dataset.metadata.input_config.bigquery_source.uri
117126
== "bq://project.dataset.table"
118127
)
119128

120-
def test_set_bigquery_uri_without_prefix(self):
121-
dataset = types.MultimodalDataset()
129+
def test_set_bigquery_uri_without_prefix(self):
130+
dataset = types.MultimodalDataset()
122131

123-
dataset.set_bigquery_uri("project.dataset.table")
132+
dataset.set_bigquery_uri("project.dataset.table")
124133

125-
assert isinstance(dataset, types.MultimodalDataset)
126-
assert (
134+
assert isinstance(dataset, types.MultimodalDataset)
135+
assert (
127136
dataset.metadata.input_config.bigquery_source.uri
128137
== "bq://project.dataset.table"
129138
)
130139

131-
def test_set_bigquery_uri_preserves_other_fields(self):
132-
dataset = types.MultimodalDataset(
140+
def test_set_bigquery_uri_preserves_other_fields(self):
141+
dataset = types.MultimodalDataset(
133142
metadata={
134143
"gemini_request_read_config": {
135144
"assembled_request_column_name": "test_column",
136145
},
137146
},
138147
)
139148

140-
dataset.set_bigquery_uri("bq://test_table")
149+
dataset.set_bigquery_uri("bq://test_table")
141150

142-
assert isinstance(dataset, types.MultimodalDataset)
143-
assert dataset.metadata.input_config.bigquery_source.uri == "bq://test_table"
144-
assert (
151+
assert isinstance(dataset, types.MultimodalDataset)
152+
assert dataset.metadata.input_config.bigquery_source.uri == "bq://test_table"
153+
assert (
145154
dataset.metadata.gemini_request_read_config.assembled_request_column_name
146155
== "test_column"
147156
)
148157

149-
def test_to_bigframes(self, mock_import_bigframes):
150-
dataset = types.MultimodalDataset()
151-
dataset.set_bigquery_uri("bq://project.dataset.table")
158+
def test_to_bigframes(self, mock_import_bigframes):
159+
dataset = types.MultimodalDataset()
160+
dataset.set_bigquery_uri("bq://project.dataset.table")
152161

153-
df = dataset.to_bigframes()
162+
df = dataset.to_bigframes()
154163

155-
assert "project.dataset.table" in df.sql
156-
mock_import_bigframes.return_value.pandas.read_gbq_table.assert_called_once_with(
164+
assert "project.dataset.table" in df.sql
165+
mock_import_bigframes.return_value.pandas.read_gbq_table.assert_called_once_with(
157166
"project.dataset.table"
158167
)
159168

169+
def test_get_batch_job_destination(self, mock_get_batch_job_unique_name):
170+
dataset = types.MultimodalDataset(
171+
name="projects/vertex-sdk-dev/locations/us-central1/datasets/12345",
172+
display_name="test_multimodal_dataset",
173+
metadata={
174+
"inputConfig": {
175+
"bigquerySource": {
176+
"uri": "bq://target_project.target_dataset.target_table"
177+
},
178+
},
179+
},
180+
)
181+
destination = dataset.get_batch_job_destination()
182+
assert (
183+
destination.vertex_dataset.display_name
184+
== "test_multimodal_dataset_batch_output_12345678901234_abcde"
185+
)
186+
assert (
187+
destination.vertex_dataset.bigquery_destination
188+
== "bq://target_project.target_dataset.target_table_batch_output_12345678901234_abcde"
189+
)
190+
191+
def test_to_batch_job_source(self):
192+
dataset = types.MultimodalDataset(
193+
name="projects/vertex-sdk-dev/locations/us-central1/datasets/12345",
194+
)
195+
source = dataset.to_batch_job_source()
196+
assert source == genai_types.BatchJobSource(
197+
vertex_dataset_name="projects/vertex-sdk-dev/locations/us-central1/datasets/12345",
198+
)
199+
160200

161201
class TestGeminiRequestReadConfig:
162202
def test_single_turn_template(self):

0 commit comments

Comments
 (0)