From 7f217877b94ec24477d9ea9d8bc1c2b4cd1866e6 Mon Sep 17 00:00:00 2001 From: NeoCodes <41864395+neo-con@users.noreply.github.com> Date: Thu, 5 Mar 2026 14:20:19 -0500 Subject: [PATCH] fix(batches): preserve order for Gemini inlined batch responses --- google/genai/batches.py | 53 ++++-- .../test_create_with_inlined_requests.py | 159 ++++++++++++++++++ google/genai/tests/conftest.py | 8 + 3 files changed, 210 insertions(+), 10 deletions(-) diff --git a/google/genai/batches.py b/google/genai/batches.py index 7a43a9aec..8fb64b60e 100644 --- a/google/genai/batches.py +++ b/google/genai/batches.py @@ -33,6 +33,7 @@ logger = logging.getLogger('google_genai.batches') +_INLINED_REQUEST_ORDER_METADATA_KEY = '_google_genai_inlined_request_order' def _AuthConfig_to_mldev( @@ -79,15 +80,37 @@ def _BatchJobDestination_from_mldev( setv(to_object, ['file_name'], getv(from_object, ['responsesFile'])) if getv(from_object, ['inlinedResponses', 'inlinedResponses']) is not None: + inlined_responses = [ + _InlinedResponse_from_mldev(item, to_object) + for item in getv(from_object, ['inlinedResponses', 'inlinedResponses']) + ] + # Backend can return inlined responses out of input order. When we have the + # SDK-injected order marker, restore the original order deterministically. + sortable = True + for inlined_response in inlined_responses: + metadata = getv(inlined_response, ['metadata']) + request_order = ( + metadata.get(_INLINED_REQUEST_ORDER_METADATA_KEY) + if isinstance(metadata, dict) + else None + ) + if request_order is None or not str(request_order).isdigit(): + sortable = False + break + if sortable: + inlined_responses.sort( + key=lambda response: int( + getv(response, ['metadata', _INLINED_REQUEST_ORDER_METADATA_KEY]) + ) + ) + for inlined_response in inlined_responses: + metadata = getv(inlined_response, ['metadata']) + if isinstance(metadata, dict): + metadata.pop(_INLINED_REQUEST_ORDER_METADATA_KEY, None) setv( to_object, ['inlined_responses'], - [ - _InlinedResponse_from_mldev(item, to_object) - for item in getv( - from_object, ['inlinedResponses', 'inlinedResponses'] - ) - ], + inlined_responses, ) if ( @@ -213,13 +236,23 @@ def _BatchJobSource_to_mldev( setv(to_object, ['fileName'], getv(from_object, ['file_name'])) if getv(from_object, ['inlined_requests']) is not None: + inlined_requests = [] + for index, inlined_request in enumerate(getv(from_object, ['inlined_requests'])): + inlined_request_object = _InlinedRequest_to_mldev( + api_client, inlined_request, to_object + ) + metadata = getv(inlined_request_object, ['metadata'], default_value={}) + if not isinstance(metadata, dict): + metadata = {} + # Reserved SDK key: always stamp deterministic order marker, even when + # caller metadata contains the same key. + metadata[_INLINED_REQUEST_ORDER_METADATA_KEY] = str(index) + setv(inlined_request_object, ['metadata'], metadata) + inlined_requests.append(inlined_request_object) setv( to_object, ['requests', 'requests'], - [ - _InlinedRequest_to_mldev(api_client, item, to_object) - for item in getv(from_object, ['inlined_requests']) - ], + inlined_requests, ) return to_object diff --git a/google/genai/tests/batches/test_create_with_inlined_requests.py b/google/genai/tests/batches/test_create_with_inlined_requests.py index f528162cd..13809fc3d 100644 --- a/google/genai/tests/batches/test_create_with_inlined_requests.py +++ b/google/genai/tests/batches/test_create_with_inlined_requests.py @@ -21,7 +21,9 @@ import os import pytest +from unittest import mock +from ... import batches as batches_module from ... import _transformers as t from ... import types from .. import pytest_helper @@ -258,6 +260,163 @@ ] +def test_inlined_requests_include_internal_order_metadata( + use_vertex, replays_prefix, http_options +): + del use_vertex, replays_prefix, http_options + request_payload = { + 'inlined_requests': [ + {'contents': [{'parts': [{'text': 'first'}], 'role': 'user'}]}, + { + 'contents': [{'parts': [{'text': 'second'}], 'role': 'user'}], + 'metadata': {'caller': 'external'}, + }, + ] + } + + converted = batches_module._BatchJobSource_to_mldev( + mock.MagicMock(), request_payload + ) + requests = converted['requests']['requests'] + key = batches_module._INLINED_REQUEST_ORDER_METADATA_KEY + + assert requests[0]['metadata'][key] == '0' + assert requests[1]['metadata'][key] == '1' + assert requests[1]['metadata']['caller'] == 'external' + + +def test_inlined_requests_internal_order_metadata_overrides_reserved_key( + use_vertex, replays_prefix, http_options +): + del use_vertex, replays_prefix, http_options + key = batches_module._INLINED_REQUEST_ORDER_METADATA_KEY + request_payload = { + 'inlined_requests': [ + { + 'contents': [{'parts': [{'text': 'first'}], 'role': 'user'}], + 'metadata': {key: '999', 'caller': 'external'}, + }, + ] + } + + converted = batches_module._BatchJobSource_to_mldev( + mock.MagicMock(), request_payload + ) + request = converted['requests']['requests'][0] + + assert request['metadata'][key] == '0' + assert request['metadata']['caller'] == 'external' + + +def test_inlined_responses_are_reordered_by_internal_order_metadata( + use_vertex, replays_prefix, http_options +): + del use_vertex, replays_prefix, http_options + key = batches_module._INLINED_REQUEST_ORDER_METADATA_KEY + response_payload = { + 'inlinedResponses': { + 'inlinedResponses': [ + { + 'metadata': {'request_key': 'two', key: '2'}, + 'response': {'candidates': []}, + }, + { + 'metadata': {'request_key': 'zero', key: '0'}, + 'response': {'candidates': []}, + }, + { + 'metadata': {'request_key': 'one', key: '1'}, + 'response': {'candidates': []}, + }, + ] + } + } + + converted = batches_module._BatchJobDestination_from_mldev(response_payload) + responses = converted['inlined_responses'] + + assert [item['metadata']['request_key'] for item in responses] == [ + 'zero', + 'one', + 'two', + ] + assert all(key not in item['metadata'] for item in responses) + + +def test_inlined_responses_keep_input_order_when_metadata_missing( + use_vertex, replays_prefix, http_options +): + del use_vertex, replays_prefix, http_options + key = batches_module._INLINED_REQUEST_ORDER_METADATA_KEY + response_payload = { + 'inlinedResponses': { + 'inlinedResponses': [ + { + 'metadata': {'request_key': 'two', key: '2'}, + 'response': {'candidates': []}, + }, + { + 'metadata': {'request_key': 'zero'}, + 'response': {'candidates': []}, + }, + { + 'metadata': {'request_key': 'one', key: '1'}, + 'response': {'candidates': []}, + }, + ] + } + } + + converted = batches_module._BatchJobDestination_from_mldev(response_payload) + responses = converted['inlined_responses'] + + assert [item['metadata']['request_key'] for item in responses] == [ + 'two', + 'zero', + 'one', + ] + assert responses[0]['metadata'][key] == '2' + assert key not in responses[1]['metadata'] + assert responses[2]['metadata'][key] == '1' + + +def test_inlined_responses_keep_input_order_when_metadata_non_numeric( + use_vertex, replays_prefix, http_options +): + del use_vertex, replays_prefix, http_options + key = batches_module._INLINED_REQUEST_ORDER_METADATA_KEY + response_payload = { + 'inlinedResponses': { + 'inlinedResponses': [ + { + 'metadata': {'request_key': 'two', key: '2'}, + 'response': {'candidates': []}, + }, + { + 'metadata': {'request_key': 'bad', key: 'not-a-number'}, + 'response': {'candidates': []}, + }, + { + 'metadata': {'request_key': 'one', key: '1'}, + 'response': {'candidates': []}, + }, + ] + } + } + + converted = batches_module._BatchJobDestination_from_mldev(response_payload) + responses = converted['inlined_responses'] + + assert [item['metadata']['request_key'] for item in responses] == [ + 'two', + 'bad', + 'one', + ] + assert responses[0]['metadata'][key] == '2' + assert responses[1]['metadata'][key] == 'not-a-number' + assert responses[2]['metadata'][key] == '1' + + @pytest.mark.asyncio async def test_async_create(client): with pytest_helper.exception_if_vertex(client, ValueError): diff --git a/google/genai/tests/conftest.py b/google/genai/tests/conftest.py index 03b813631..6a74d6f73 100644 --- a/google/genai/tests/conftest.py +++ b/google/genai/tests/conftest.py @@ -92,6 +92,14 @@ def client(use_vertex, replays_prefix, http_options, request): Assert an exception if the test is not supported in an API.""") replay_id = _get_replay_id(use_vertex, replays_prefix) + if mode in ['replay', 'tap'] and not use_vertex: + # Replay mode should not require a real API key, but client init still + # validates key presence on the mldev path. + if not os.environ.get('GOOGLE_API_KEY') and not os.environ.get( + 'GEMINI_API_KEY' + ): + os.environ['GOOGLE_API_KEY'] = 'dummy-api-key' + if mode == 'tap': mode = 'replay'