Skip to content

Commit f138162

Browse files
cleop-googlecopybara-github
authored andcommitted
feat: GenAI SDK client(multimodal) - Accept gemini_request_read_config instead of template_config in all functions.
PiperOrigin-RevId: 888570221
1 parent 4f0fdfe commit f138162

File tree

4 files changed

+223
-206
lines changed

4 files changed

+223
-206
lines changed

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

Lines changed: 24 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -41,16 +41,18 @@ def test_assemble_dataset(client):
4141
def test_assemble_dataset_public(client):
4242
bigquery_destination = client.datasets.assemble(
4343
name=DATASET,
44-
template_config=types.GeminiTemplateConfig(
45-
gemini_example=types.GeminiExample(
46-
model="gemini-1.5-flash",
47-
contents=[
48-
{
49-
"role": "user",
50-
"parts": [{"text": "What is the capital of {name}?"}],
51-
}
52-
],
53-
),
44+
gemini_request_read_config=types.GeminiRequestReadConfig(
45+
template_config=types.GeminiTemplateConfig(
46+
gemini_example=types.GeminiExample(
47+
model="gemini-1.5-flash",
48+
contents=[
49+
{
50+
"role": "user",
51+
"parts": [{"text": "What is the capital of {name}?"}],
52+
}
53+
],
54+
),
55+
)
5456
),
5557
)
5658
assert bigquery_destination.startswith(f"bq://{BIGQUERY_TABLE_NAME}")
@@ -81,16 +83,18 @@ async def test_assemble_dataset_async(client):
8183
async def test_assemble_dataset_public_async(client):
8284
bigquery_destination = await client.aio.datasets.assemble(
8385
name=DATASET,
84-
template_config=types.GeminiTemplateConfig(
85-
gemini_example=types.GeminiExample(
86-
model="gemini-1.5-flash",
87-
contents=[
88-
{
89-
"role": "user",
90-
"parts": [{"text": "What is the capital of {name}?"}],
91-
}
92-
],
93-
),
86+
gemini_request_read_config=types.GeminiRequestReadConfig(
87+
template_config=types.GeminiTemplateConfig(
88+
gemini_example=types.GeminiExample(
89+
model="gemini-1.5-flash",
90+
contents=[
91+
{
92+
"role": "user",
93+
"parts": [{"text": "What is the capital of {name}?"}],
94+
}
95+
],
96+
),
97+
)
9498
),
9599
)
96100
assert bigquery_destination.startswith(f"bq://{BIGQUERY_TABLE_NAME}")

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

Lines changed: 112 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -52,15 +52,17 @@ def test_assess_tuning_resources(client):
5252
response = client.datasets.assess_tuning_resources(
5353
dataset_name=DATASET,
5454
model_name="gemini-2.5-flash-001",
55-
template_config=types.GeminiTemplateConfig(
56-
gemini_example=types.GeminiExample(
57-
contents=[
58-
{
59-
"role": "user",
60-
"parts": [{"text": "What is the capital of {name}?"}],
61-
}
62-
],
63-
),
55+
gemini_request_read_config=types.GeminiRequestReadConfig(
56+
template_config=types.GeminiTemplateConfig(
57+
gemini_example=types.GeminiExample(
58+
contents=[
59+
{
60+
"role": "user",
61+
"parts": [{"text": "What is the capital of {name}?"}],
62+
}
63+
],
64+
),
65+
)
6466
),
6567
)
6668
assert isinstance(response, types.TuningResourceUsageAssessmentResult)
@@ -71,19 +73,21 @@ def test_assess_tuning_validity(client):
7173
dataset_name=DATASET,
7274
dataset_usage="SFT_VALIDATION",
7375
model_name="gemini-2.5-flash-001",
74-
template_config=types.GeminiTemplateConfig(
75-
gemini_example=types.GeminiExample(
76-
contents=[
77-
{
78-
"role": "user",
79-
"parts": [{"text": "What is the capital of {name}?"}],
80-
},
81-
{
82-
"role": "model",
83-
"parts": [{"text": "{capital}"}],
84-
},
85-
],
86-
),
76+
gemini_request_read_config=types.GeminiRequestReadConfig(
77+
template_config=types.GeminiTemplateConfig(
78+
gemini_example=types.GeminiExample(
79+
contents=[
80+
{
81+
"role": "user",
82+
"parts": [{"text": "What is the capital of {name}?"}],
83+
},
84+
{
85+
"role": "model",
86+
"parts": [{"text": "{capital}"}],
87+
},
88+
],
89+
),
90+
)
8791
),
8892
)
8993
assert isinstance(response, types.TuningValidationAssessmentResult)
@@ -93,19 +97,21 @@ def test_assess_batch_prediction_resources(client):
9397
response = client.datasets.assess_batch_prediction_resources(
9498
dataset_name=DATASET,
9599
model_name="gemini-2.5-flash-001",
96-
template_config=types.GeminiTemplateConfig(
97-
gemini_example=types.GeminiExample(
98-
contents=[
99-
{
100-
"role": "user",
101-
"parts": [{"text": "What is the capital of {name}?"}],
102-
},
103-
{
104-
"role": "model",
105-
"parts": [{"text": "{capital}"}],
106-
},
107-
],
108-
),
100+
gemini_request_read_config=types.GeminiRequestReadConfig(
101+
template_config=types.GeminiTemplateConfig(
102+
gemini_example=types.GeminiExample(
103+
contents=[
104+
{
105+
"role": "user",
106+
"parts": [{"text": "What is the capital of {name}?"}],
107+
},
108+
{
109+
"role": "model",
110+
"parts": [{"text": "{capital}"}],
111+
},
112+
],
113+
),
114+
)
109115
),
110116
)
111117
assert isinstance(response, types.BatchPredictionResourceUsageAssessmentResult)
@@ -115,19 +121,21 @@ def test_assess_batch_prediction_validity(client):
115121
response = client.datasets.assess_batch_prediction_validity(
116122
dataset_name=DATASET,
117123
model_name="gemini-2.5-flash-001",
118-
template_config=types.GeminiTemplateConfig(
119-
gemini_example=types.GeminiExample(
120-
contents=[
121-
{
122-
"role": "user",
123-
"parts": [{"text": "What is the capital of {name}?"}],
124-
},
125-
{
126-
"role": "model",
127-
"parts": [{"text": "{capital}"}],
128-
},
129-
],
130-
),
124+
gemini_request_read_config=types.GeminiRequestReadConfig(
125+
template_config=types.GeminiTemplateConfig(
126+
gemini_example=types.GeminiExample(
127+
contents=[
128+
{
129+
"role": "user",
130+
"parts": [{"text": "What is the capital of {name}?"}],
131+
},
132+
{
133+
"role": "model",
134+
"parts": [{"text": "{capital}"}],
135+
},
136+
],
137+
),
138+
)
131139
),
132140
)
133141
assert isinstance(response, types.BatchPredictionValidationAssessmentResult)
@@ -169,15 +177,17 @@ async def test_assess_tuning_resources_async(client):
169177
response = await client.aio.datasets.assess_tuning_resources(
170178
dataset_name=DATASET,
171179
model_name="gemini-2.5-flash-001",
172-
template_config=types.GeminiTemplateConfig(
173-
gemini_example=types.GeminiExample(
174-
contents=[
175-
{
176-
"role": "user",
177-
"parts": [{"text": "What is the capital of {name}?"}],
178-
}
179-
],
180-
),
180+
gemini_request_read_config=types.GeminiRequestReadConfig(
181+
template_config=types.GeminiTemplateConfig(
182+
gemini_example=types.GeminiExample(
183+
contents=[
184+
{
185+
"role": "user",
186+
"parts": [{"text": "What is the capital of {name}?"}],
187+
}
188+
],
189+
),
190+
)
181191
),
182192
)
183193
assert isinstance(response, types.TuningResourceUsageAssessmentResult)
@@ -189,19 +199,21 @@ async def test_assess_tuning_validity_async(client):
189199
dataset_name=DATASET,
190200
dataset_usage="SFT_VALIDATION",
191201
model_name="gemini-2.5-flash-001",
192-
template_config=types.GeminiTemplateConfig(
193-
gemini_example=types.GeminiExample(
194-
contents=[
195-
{
196-
"role": "user",
197-
"parts": [{"text": "What is the capital of {name}?"}],
198-
},
199-
{
200-
"role": "model",
201-
"parts": [{"text": "{capital}"}],
202-
},
203-
],
204-
),
202+
gemini_request_read_config=types.GeminiRequestReadConfig(
203+
template_config=types.GeminiTemplateConfig(
204+
gemini_example=types.GeminiExample(
205+
contents=[
206+
{
207+
"role": "user",
208+
"parts": [{"text": "What is the capital of {name}?"}],
209+
},
210+
{
211+
"role": "model",
212+
"parts": [{"text": "{capital}"}],
213+
},
214+
],
215+
),
216+
)
205217
),
206218
)
207219
assert isinstance(response, types.TuningValidationAssessmentResult)
@@ -212,19 +224,21 @@ async def test_assess_batch_prediction_resources_async(client):
212224
response = await client.aio.datasets.assess_batch_prediction_resources(
213225
dataset_name=DATASET,
214226
model_name="gemini-2.5-flash-001",
215-
template_config=types.GeminiTemplateConfig(
216-
gemini_example=types.GeminiExample(
217-
contents=[
218-
{
219-
"role": "user",
220-
"parts": [{"text": "What is the capital of {name}?"}],
221-
},
222-
{
223-
"role": "model",
224-
"parts": [{"text": "{capital}"}],
225-
},
226-
],
227-
),
227+
gemini_request_read_config=types.GeminiRequestReadConfig(
228+
template_config=types.GeminiTemplateConfig(
229+
gemini_example=types.GeminiExample(
230+
contents=[
231+
{
232+
"role": "user",
233+
"parts": [{"text": "What is the capital of {name}?"}],
234+
},
235+
{
236+
"role": "model",
237+
"parts": [{"text": "{capital}"}],
238+
},
239+
],
240+
),
241+
)
228242
),
229243
)
230244
assert isinstance(response, types.BatchPredictionResourceUsageAssessmentResult)
@@ -235,19 +249,21 @@ async def test_assess_batch_prediction_validity_async(client):
235249
response = await client.aio.datasets.assess_batch_prediction_validity(
236250
dataset_name=DATASET,
237251
model_name="gemini-2.5-flash-001",
238-
template_config=types.GeminiTemplateConfig(
239-
gemini_example=types.GeminiExample(
240-
contents=[
241-
{
242-
"role": "user",
243-
"parts": [{"text": "What is the capital of {name}?"}],
244-
},
245-
{
246-
"role": "model",
247-
"parts": [{"text": "{capital}"}],
248-
},
249-
],
250-
),
252+
gemini_request_read_config=types.GeminiRequestReadConfig(
253+
template_config=types.GeminiTemplateConfig(
254+
gemini_example=types.GeminiExample(
255+
contents=[
256+
{
257+
"role": "user",
258+
"parts": [{"text": "What is the capital of {name}?"}],
259+
},
260+
{
261+
"role": "model",
262+
"parts": [{"text": "{capital}"}],
263+
},
264+
],
265+
),
266+
)
251267
),
252268
)
253269
assert isinstance(response, types.BatchPredictionValidationAssessmentResult)

0 commit comments

Comments
 (0)