Skip to content

Commit a4ea2dd

Browse files
committed
fix: parse mldev batch input config source
1 parent 84822a1 commit a4ea2dd

2 files changed

Lines changed: 73 additions & 0 deletions

File tree

google/genai/batches.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,27 @@ def _BatchJobSource_from_vertex(
238238
return to_object
239239

240240

241+
def _BatchJobSource_from_mldev(
242+
from_object: Union[dict[str, Any], object],
243+
parent_object: Optional[dict[str, Any]] = None,
244+
) -> dict[str, Any]:
245+
to_object: dict[str, Any] = {}
246+
if getv(from_object, ['fileName']) is not None:
247+
setv(to_object, ['file_name'], getv(from_object, ['fileName']))
248+
249+
if getv(from_object, ['requests', 'requests']) is not None:
250+
setv(
251+
to_object,
252+
['inlined_requests'],
253+
[
254+
_InlinedRequest_from_mldev(item, to_object)
255+
for item in getv(from_object, ['requests', 'requests'])
256+
],
257+
)
258+
259+
return to_object
260+
261+
241262
def _BatchJobSource_to_mldev(
242263
api_client: BaseApiClient,
243264
from_object: Union[dict[str, Any], object],
@@ -366,6 +387,15 @@ def _BatchJob_from_mldev(
366387
if getv(from_object, ['metadata', 'model']) is not None:
367388
setv(to_object, ['model'], getv(from_object, ['metadata', 'model']))
368389

390+
if getv(from_object, ['metadata', 'inputConfig']) is not None:
391+
setv(
392+
to_object,
393+
['src'],
394+
_BatchJobSource_from_mldev(
395+
getv(from_object, ['metadata', 'inputConfig']), to_object
396+
),
397+
)
398+
369399
if getv(from_object, ['metadata', 'output']) is not None:
370400
setv(
371401
to_object,
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
16+
"""Tests for batch response converters."""
17+
18+
from ... import batches
19+
from ... import types
20+
21+
22+
def test_mldev_batch_job_parses_input_config_src():
23+
response = {
24+
'name': 'batches/test-batch',
25+
'metadata': {
26+
'inputConfig': {
27+
'fileName': 'files/input-file',
28+
},
29+
'output': {
30+
'responsesFile': 'files/output-file',
31+
},
32+
},
33+
}
34+
35+
parsed = types.BatchJob._from_response(
36+
response=batches._BatchJob_from_mldev(response),
37+
kwargs={},
38+
)
39+
40+
assert parsed.src is not None
41+
assert parsed.src.file_name == 'files/input-file'
42+
assert parsed.dest is not None
43+
assert parsed.dest.file_name == 'files/output-file'

0 commit comments

Comments
 (0)