google--langextract/tests/openai_batch_test.py

# Copyright 2025 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for OpenAI Batch API helper."""

# pylint: disable=too-few-public-methods

from __future__ import annotations

import json
import types as py_types
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized

from langextract.core import exceptions
from langextract.providers import openai_batch


class _FakeFiles:

  def __init__(self):
    self.created = []
    self.deleted = []
    self._content_by_id = {}

  def create(self, *, file, purpose):
    self.created.append({'file': file, 'purpose': purpose})
    return py_types.SimpleNamespace(id=f'file-{len(self.created)}')

  def content(self, file_id):
    return py_types.SimpleNamespace(text=self._content_by_id[file_id])

  def delete(self, file_id):
    self.deleted.append(file_id)

  def set_content(self, file_id: str, text: str) -> None:
    self._content_by_id[file_id] = text


class _FakeBatches:

  def __init__(self):
    self.created = []
    self.cancelled = []
    self._retrieve_queue = []

  def create(self, **kwargs):
    self.created.append(kwargs)
    return py_types.SimpleNamespace(id=f'batch-{len(self.created)}')

  def retrieve(self, _batch_id):
    if not self._retrieve_queue:
      raise RuntimeError('retrieve queue empty')
    return self._retrieve_queue.pop(0)

  def cancel(self, batch_id):
    self.cancelled.append(batch_id)

  def push_retrieve(self, obj):
    self._retrieve_queue.append(obj)


class _FakeClient:

  def __init__(self):
    self.files = _FakeFiles()
    self.batches = _FakeBatches()


class _StatusError(Exception):

  def __init__(self, status_code: int):
    super().__init__(f'Error code: {status_code}')
    self.status_code = status_code


def _make_output_line(custom_id: str, content: str) -> str:
  obj = {
      'custom_id': custom_id,
      'response': {
          'status_code': 200,
          'body': {
              'choices': [
                  {'message': {'content': content}},
              ]
          },
      },
      'error': None,
  }
  return json.dumps(obj)


def _make_error_line(
    custom_id: str, message: str, status_code: int = 400
) -> str:
  obj = {
      'custom_id': custom_id,
      'response': {
          'status_code': status_code,
          'body': {'error': {'message': message}},
      },
      'error': {'message': message},
  }
  return json.dumps(obj)


def _make_refusal_line(custom_id: str, refusal: str) -> str:
  obj = {
      'custom_id': custom_id,
      'response': {
          'status_code': 200,
          'body': {
              'choices': [
                  {'message': {'content': None, 'refusal': refusal}},
              ]
          },
      },
      'error': None,
  }
  return json.dumps(obj)


class OpenAIBatchConfigTest(absltest.TestCase):

  def test_default_timeout_covers_completion_window(self):
    cfg = openai_batch.BatchConfig(enabled=True)

    self.assertGreater(cfg.timeout, 24 * 60 * 60)

  def test_from_dict_accepts_batch_config(self):
    cfg = openai_batch.BatchConfig(enabled=True, threshold=1)

    self.assertIs(openai_batch.BatchConfig.from_dict(cfg), cfg)

  def test_from_dict_accepts_boolean_shorthand(self):
    enabled = openai_batch.BatchConfig.from_dict(True)
    disabled = openai_batch.BatchConfig.from_dict(False)

    self.assertTrue(enabled.enabled)
    self.assertFalse(disabled.enabled)

  def test_from_dict_rejects_invalid_type(self):
    with self.assertRaisesRegex(TypeError, 'batch must be a mapping'):
      openai_batch.BatchConfig.from_dict([])

  def test_infer_batch_rejects_invalid_batch_size(self):
    cfg = openai_batch.BatchConfig(enabled=True, threshold=1)

    with self.assertRaisesRegex(
        exceptions.InferenceConfigError, 'batch_size must be > 0'
    ):
      openai_batch.infer_batch(
          client=_FakeClient(),
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
          batch_size=0,
      )


class OpenAIBatchHelperTest(parameterized.TestCase):

  @mock.patch(
      'langextract.providers.openai_batch.time.sleep', return_value=None
  )
  def test_orders_results_by_custom_id(self, _mock_sleep):
    client = _FakeClient()

    client.batches.push_retrieve(py_types.SimpleNamespace(status='in_progress'))
    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )

    out = '\n'.join([
        _make_output_line('idx-000001', 'B'),
        _make_output_line('idx-000000', 'A'),
    ])
    client.files.set_content('out-1', out)

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    res = openai_batch.infer_batch(
        client=client,
        model_id='gpt-test',
        prompts=['p0', 'p1'],
        cfg=cfg,
        request_builder=lambda p: {
            'model': 'gpt-test',
            'messages': [{'role': 'user', 'content': p}],
        },
    )

    self.assertEqual(res, ['A', 'B'])

  @mock.patch(
      'langextract.providers.openai_batch.time.sleep', return_value=None
  )
  def test_splits_jobs_by_batch_size(self, _mock_sleep):
    client = _FakeClient()

    for job_idx in range(3):
      client.batches.push_retrieve(
          py_types.SimpleNamespace(status='in_progress')
      )
      client.batches.push_retrieve(
          py_types.SimpleNamespace(
              status='completed', output_file_id=f'out-{job_idx}'
          )
      )

    client.files.set_content(
        'out-0',
        '\n'.join([
            _make_output_line('idx-000000', '0'),
            _make_output_line('idx-000001', '1'),
        ]),
    )
    client.files.set_content(
        'out-1',
        '\n'.join([
            _make_output_line('idx-000002', '2'),
            _make_output_line('idx-000003', '3'),
        ]),
    )
    client.files.set_content(
        'out-2',
        _make_output_line('idx-000004', '4'),
    )

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
        max_requests_per_job=100,
    )

    prompts = ['p0', 'p1', 'p2', 'p3', 'p4']
    res = openai_batch.infer_batch(
        client=client,
        model_id='gpt-test',
        prompts=prompts,
        cfg=cfg,
        request_builder=lambda p: {
            'model': 'gpt-test',
            'messages': [{'role': 'user', 'content': p}],
        },
        batch_size=2,
    )

    self.assertEqual(res, ['0', '1', '2', '3', '4'])
    self.assertLen(client.batches.created, 3)
    self.assertLen(client.files.created, 3)

  def test_metadata_and_job_create_hook_are_used(self):
    client = _FakeClient()
    created_jobs = []

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.set_content('out-1', _make_output_line('idx-000000', 'ok'))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
        metadata={'purpose': 'test'},
        on_job_create=created_jobs.append,
    )

    res = openai_batch.infer_batch(
        client=client,
        model_id='gpt-test',
        prompts=['p0'],
        cfg=cfg,
        request_builder=lambda p: {
            'model': 'gpt-test',
            'messages': [{'role': 'user', 'content': p}],
        },
    )

    self.assertEqual(res, ['ok'])
    self.assertEqual(client.batches.created[0]['metadata'], {'purpose': 'test'})
    self.assertLen(created_jobs, 1)
    self.assertEqual(created_jobs[0].id, 'batch-1')

  def test_item_error_raises(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )

    obj = {
        'custom_id': 'idx-000000',
        'error': {'message': 'boom'},
        'response': None,
    }
    client.files.set_content('out-1', json.dumps(obj))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'per-item errors'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  def test_job_create_failure_deletes_uploaded_input_file(self):
    client = _FakeClient()
    client.batches.create = mock.Mock(side_effect=RuntimeError('boom'))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'job create failed'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

    self.assertEqual(client.files.deleted, ['file-1'])

  def test_failed_job_reports_errors_field(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(
            status='failed',
            errors={'data': [{'message': 'invalid request'}]},
        )
    )

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'invalid request'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  @parameterized.named_parameters(
      dict(testcase_name='failed', status='failed'),
      dict(testcase_name='expired', status='expired'),
      dict(testcase_name='cancelled', status='cancelled'),
  )
  def test_terminal_job_status_reports_status(self, status):
    client = _FakeClient()

    client.batches.push_retrieve(py_types.SimpleNamespace(status=status))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, f'status={status}'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  def test_completed_job_missing_output_reports_custom_id(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.set_content('out-1', _make_output_line('idx-000000', 'ok'))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'custom_id=idx-000001'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0', 'p1'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  def test_completed_job_without_output_or_error_file_raises(self):
    client = _FakeClient()

    client.batches.push_retrieve(py_types.SimpleNamespace(status='completed'))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'no output_file_id or error_file_id'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  def test_unexpected_custom_id_logs_warning(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.set_content('out-1', _make_output_line('bad-id', 'ok'))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertLogs(level='WARNING') as logs:
      with self.assertRaisesRegex(
          exceptions.InferenceRuntimeError, 'custom_id=idx-000000'
      ):
        openai_batch.infer_batch(
            client=client,
            model_id='gpt-test',
            prompts=['p0'],
            cfg=cfg,
            request_builder=lambda p: {
                'model': 'gpt-test',
                'messages': [{'role': 'user', 'content': p}],
            },
        )

    self.assertIn('unexpected custom_id', '\n'.join(logs.output))

  @mock.patch(
      'langextract.providers.openai_batch.time.sleep', return_value=None
  )
  def test_output_download_retries_transient_forbidden(self, mock_sleep):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.content = mock.Mock(
        side_effect=[
            _StatusError(403),
            py_types.SimpleNamespace(
                text=_make_output_line('idx-000000', 'ok')
            ),
        ]
    )

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    res = openai_batch.infer_batch(
        client=client,
        model_id='gpt-test',
        prompts=['p0'],
        cfg=cfg,
        request_builder=lambda p: {
            'model': 'gpt-test',
            'messages': [{'role': 'user', 'content': p}],
        },
    )

    self.assertEqual(res, ['ok'])
    self.assertEqual(client.files.content.call_count, 2)
    mock_sleep.assert_called_once_with(1)

  @mock.patch(
      'langextract.providers.openai_batch.time.sleep', return_value=None
  )
  def test_output_download_forbidden_mentions_files_permission(
      self, mock_sleep
  ):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.content = mock.Mock(side_effect=_StatusError(403))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'Files Read permission'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )
    self.assertEqual(client.files.content.call_count, 3)
    self.assertEqual(mock_sleep.call_count, 2)

  def test_completed_job_reads_error_file(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(
            status='completed',
            output_file_id='out-1',
            error_file_id='err-1',
        )
    )
    client.files.set_content('out-1', _make_output_line('idx-000000', 'ok'))
    client.files.set_content(
        'err-1', _make_error_line('idx-000001', 'rate limited', 429)
    )

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'rate limited'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0', 'p1'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  def test_non_2xx_response_raises_item_error(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.set_content(
        'out-1', _make_error_line('idx-000000', 'rate limited', 429)
    )

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'status_code=429'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  def test_refusal_message_is_reported(self):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.set_content(
        'out-1', _make_refusal_line('idx-000000', 'cannot comply')
    )

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(
        exceptions.InferenceRuntimeError, 'cannot comply'
    ):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

  @mock.patch('langextract.providers.openai_batch.time.time')
  def test_timeout_cancels_job(self, mock_time):
    # The second clock read crosses the timeout immediately, without sleeping.
    mock_time.side_effect = [0, 10]
    client = _FakeClient()

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        completion_window='24h',
        poll_interval=1,
        timeout=5,
    )

    with self.assertRaisesRegex(exceptions.InferenceRuntimeError, 'timed out'):
      openai_batch.infer_batch(
          client=client,
          model_id='gpt-test',
          prompts=['p0'],
          cfg=cfg,
          request_builder=lambda p: {
              'model': 'gpt-test',
              'messages': [{'role': 'user', 'content': p}],
          },
      )

    self.assertEqual(client.batches.cancelled, ['batch-1'])

  @mock.patch(
      'langextract.providers.openai_batch.time.sleep', return_value=None
  )
  def test_default_completion_window_is_sent(self, _mock_sleep):
    client = _FakeClient()

    client.batches.push_retrieve(
        py_types.SimpleNamespace(status='completed', output_file_id='out-1')
    )
    client.files.set_content('out-1', _make_output_line('idx-000000', 'ok'))

    cfg = openai_batch.BatchConfig(
        enabled=True,
        threshold=1,
        poll_interval=1,
        timeout=5,
    )

    res = openai_batch.infer_batch(
        client=client,
        model_id='gpt-test',
        prompts=['p0'],
        cfg=cfg,
        request_builder=lambda p: {
            'model': 'gpt-test',
            'messages': [{'role': 'user', 'content': p}],
        },
    )

    self.assertEqual(res, ['ok'])
    self.assertLen(client.batches.created, 1)
    self.assertEqual(client.batches.created[0]['completion_window'], '24h')

  def test_unsupported_completion_window_raises(self):
    with self.assertRaisesRegex(ValueError, 'completion_window'):
      openai_batch.BatchConfig(enabled=True, completion_window='1h')


if __name__ == '__main__':
  absltest.main()