Fix ordering of Gemini inlined batch responses

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

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):

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'