first vommit

2026-07-04 14:31:50 +02:00
commit a04e68bc95
@@ -0,0 +1,17 @@
+# 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 the Google GenAI SDK's batches module."""
@@ -0,0 +1,77 @@
+# 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 batches.cancel()."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+# Vertex AI batch job name.
+_BATCH_JOB_NAME = '6339625664542408704'
+_BATCH_JOB_FULL_RESOURCE_NAME = (
+    'projects/964831358985/locations/us-central1/'
+    f'batchPredictionJobs/{_BATCH_JOB_NAME}'
+)
+# MLDev batch operation name.
+_MLDEV_BATCH_OPERATION_NAME = 'batches/0yew7plxupyybd7appsrq5vw7w0lp3l79lab'
+_INVALID_BATCH_JOB_NAME = 'invalid_name'
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_cancel_batch_job',
+        parameters=types._CancelBatchJobParameters(
+            name=_BATCH_JOB_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+        skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_cancel_batch_operation',
+        parameters=types._CancelBatchJobParameters(
+            name=_MLDEV_BATCH_OPERATION_NAME,
+        ),
+        exception_if_vertex='Invalid batch job name',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_cancel_batch_job_with_invalid_name',
+        parameters=types._CancelBatchJobParameters(
+            name=_INVALID_BATCH_JOB_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+        exception_if_vertex='Invalid batch job name',
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='batches.cancel',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async_cancel(client):
+  if client.vertexai:
+    name = _BATCH_JOB_NAME
+  else:
+    name = _MLDEV_BATCH_OPERATION_NAME
+  await client.aio.batches.cancel(name=name)
@@ -0,0 +1,123 @@
+# 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 batches.create() with invalid source and destinations."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+_GEMINI_MODEL = 'gemini-2.5-flash'
+_EMBEDDING_MODEL = 'gemini-embedding-001'
+_DISPLAY_NAME = 'test_batch'
+
+_GENERATE_CONTENT_BQ_OUTPUT_PREFIX = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
+)
+
+_EMBEDDING_BQ_INPUT_FILE = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_requests'
+)
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_union_with_invalid_src',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL,
+            src='invalid_src',
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': _GENERATE_CONTENT_BQ_OUTPUT_PREFIX,
+            },
+        ),
+        exception_if_mldev='Unsupported source',
+        exception_if_vertex='Unsupported source',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_union_with_invalid_dest',
+        parameters=types._CreateBatchJobParameters(
+            model=_EMBEDDING_MODEL,
+            src=_EMBEDDING_BQ_INPUT_FILE,
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': 'invalid_dest',
+            },
+        ),
+        exception_if_mldev='supported',
+        exception_if_vertex='Unsupported destination',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_create_with_webhook_config',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL,
+            src={
+                'inlined_requests': [
+                    {
+                        'contents': [
+                            {'parts': [{'text': 'say hello'}], 'role': 'user'}
+                        ]
+                    },
+                ]
+            },
+            config=types.CreateBatchJobConfig(
+                display_name=_DISPLAY_NAME,
+                webhook_config=types.WebhookConfig(
+                    uris=['https://example.com/webhook'],
+                    user_metadata={'batch_id': '123'},
+                ),
+            ),
+        ),
+        exception_if_vertex='supported',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_create_with_webhook_config_dict',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL,
+            src={
+                'inlined_requests': [
+                    {
+                        'contents': [
+                            {'parts': [{'text': 'say hello'}], 'role': 'user'}
+                        ]
+                    },
+                ]
+            },
+            config={
+                'display_name': _DISPLAY_NAME,
+                'webhook_config': {
+                    'uris': ['https://example.com/webhook'],
+                },
+            },
+        ),
+        exception_if_vertex='supported',
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create',
+        test_table=test_table,
+    ),
+]
@@ -0,0 +1,113 @@
+# 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 batches.create() with BigQuery source."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+_GEMINI_MODEL = 'gemini-2.5-flash'
+_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
+_EMBEDDING_MODEL = 'gemini-embedding-001'
+_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
+_DISPLAY_NAME = 'test_batch'
+_GENERATE_CONTENT_BQ_INPUT_FILE = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_requests'
+)
+_GENERATE_CONTENT_BQ_OUTPUT_PREFIX = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
+)
+
+_EMBEDDING_BQ_INPUT_FILE = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_requests'
+)
+_EMBEDDING_BQ_OUTPUT_PREFIX = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.embedding_output'
+)
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_union_generate_content_with_bigquery',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL_FULL_NAME,
+            src=_GENERATE_CONTENT_BQ_INPUT_FILE,
+        ),
+        exception_if_mldev='not supported in Gemini API',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_union_embedding_with_bigquery',
+        parameters=types._CreateBatchJobParameters(
+            model='publishers/google/models/text-embedding-004',
+            src=_EMBEDDING_BQ_INPUT_FILE,
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': _EMBEDDING_BQ_OUTPUT_PREFIX,
+            },
+        ),
+        exception_if_mldev='not supported in Gemini API',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_generate_content_with_bigquery',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL_FULL_NAME,
+            src={
+                'bigquery_uri': _GENERATE_CONTENT_BQ_INPUT_FILE,
+                'format': 'bigquery',
+            },
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': {
+                    'bigquery_uri': _GENERATE_CONTENT_BQ_OUTPUT_PREFIX,
+                    'format': 'bigquery',
+                },
+            },
+        ),
+        exception_if_mldev='one of',
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create',
+        test_table=test_table,
+    ),
+]
+
+
+@pytest.mark.asyncio
+async def test_async_create(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    batch_job = await client.aio.batches.create(
+        model=_GEMINI_MODEL,
+        src=_GENERATE_CONTENT_BQ_INPUT_FILE,
+    )
+
+    assert batch_job.name.startswith('projects/')
+    assert (
+        batch_job.model == _GEMINI_MODEL_FULL_NAME
+    )  # Converted to Vertex full name.
+    assert batch_job.src.bigquery_uri == _GENERATE_CONTENT_BQ_INPUT_FILE
+    assert batch_job.src.format == 'bigquery'
@@ -0,0 +1,82 @@
+# 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 batches.create() with file source."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+_GEMINI_MODEL = 'gemini-2.5-flash'
+_DISPLAY_NAME = 'test_batch'
+
+_MLDEV_GEMINI_MODEL = 'gemini-2.5-flash'
+_FILE_NAME = 'files/s0pa54alni6w'
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_union_generate_content_with_file',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'file_name': _FILE_NAME},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_generate_content_with_file',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'file_name': _FILE_NAME},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported',
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create',
+        test_table=test_table,
+    ),
+]
+
+
+@pytest.mark.asyncio
+async def test_async_create(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    batch_job = await client.aio.batches.create(
+        model=_MLDEV_GEMINI_MODEL,
+        src={'file_name': _FILE_NAME},
+        config={
+            'display_name': _DISPLAY_NAME,
+        },
+    )
+    assert batch_job.name.startswith('batches/')
+    assert (
+        batch_job.model == 'models/' + _MLDEV_GEMINI_MODEL
+    )  # Converted to Gemini full name.
@@ -0,0 +1,125 @@
+# 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 batches.create() with Google Cloud Storage source."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+_GEMINI_MODEL = 'gemini-2.5-flash'
+_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
+_EMBEDDING_MODEL = 'gemini-embedding-001'
+_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/gemini-embedding-001'
+_DISPLAY_NAME = 'test_batch'
+
+_GENERATE_CONTENT_GCS_INPUT_FILE = (
+    'gs://unified-genai-tests/batches/input/generate_content_requests.jsonl'
+)
+_GENERATE_CONTENT_GCS_OUTPUT_PREFIX = 'gs://unified-genai-tests/batches/output'
+
+_EMBEDDING_GCS_INPUT_FILE = (
+    'gs://unified-genai-tests/batches/input/embedding_requests.jsonl'
+)
+_EMBEDDING_GCS_OUTPUT_PREFIX = 'gs://unified-genai-tests/batches/output'
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_union_generate_content_with_gcs',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL,
+            src=_GENERATE_CONTENT_GCS_INPUT_FILE,
+        ),
+        exception_if_mldev='not supported in Gemini API',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_union_embedding_with_gcs',
+        parameters=types._CreateBatchJobParameters(
+            model='text-embedding-004',
+            src=_EMBEDDING_GCS_INPUT_FILE,
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': _EMBEDDING_GCS_OUTPUT_PREFIX,
+            },
+        ),
+        exception_if_mldev='not supported in Gemini API',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_union_with_http_options',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL,
+            src=_GENERATE_CONTENT_GCS_INPUT_FILE,
+            config={
+                'http_options': {
+                    'api_version': 'v1',
+                    'headers': {'test': 'headers'},
+                },
+            },
+        ),
+        exception_if_mldev='not supported in Gemini API',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_generate_content_with_gcs',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL,
+            src={
+                'gcs_uri': [_GENERATE_CONTENT_GCS_INPUT_FILE],
+                'format': 'jsonl',
+            },
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': {
+                    'gcs_uri': _GENERATE_CONTENT_GCS_OUTPUT_PREFIX,
+                    'format': 'jsonl',
+                },
+            },
+        ),
+        exception_if_mldev='one of',
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create',
+        test_table=test_table,
+    ),
+]
+
+
+@pytest.mark.asyncio
+async def test_async_create(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    batch_job = await client.aio.batches.create(
+        model=_GEMINI_MODEL,
+        src=_GENERATE_CONTENT_GCS_INPUT_FILE,
+    )
+
+    assert batch_job.name.startswith('projects/')
+    assert (
+        batch_job.model == _GEMINI_MODEL_FULL_NAME
+    )  # Converted to Vertex full name.
+    assert batch_job.src.gcs_uri == [_GENERATE_CONTENT_GCS_INPUT_FILE]
+    assert batch_job.src.format == 'jsonl'
@@ -0,0 +1,271 @@
+# 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 batches.create() with inlined requests."""
+import base64
+import copy
+import datetime
+import os
+
+import pytest
+
+from ... import _transformers as t
+from ... import types
+from .. import pytest_helper
+
+_GEMINI_MODEL = 'gemini-2.5-flash'
+_DISPLAY_NAME = 'test_batch'
+
+_MLDEV_GEMINI_MODEL = 'gemini-2.5-flash'
+
+_SAFETY_SETTINGS = [
+    {
+        'category': 'HARM_CATEGORY_HATE_SPEECH',
+        'threshold': 'BLOCK_ONLY_HIGH',
+    },
+    {
+        'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',
+        'threshold': 'BLOCK_LOW_AND_ABOVE',
+    },
+]
+
+_INLINED_REQUESTS = [
+    {
+        'contents': [{
+            'parts': [{
+                'text': 'what is the number after 1? return just the number.',
+            }],
+            'role': 'user',
+        }],
+        'metadata': {
+            'key': 'request-1',
+        },
+        'config': {
+            'safety_settings': _SAFETY_SETTINGS,
+        },
+    },
+    {
+        'contents': [{
+            'parts': [{
+                'text': 'what is the number after 2? return just the number.',
+            }],
+            'role': 'user',
+        }],
+        'metadata': {
+            'key': 'request-2',
+        },
+        'config': {
+            'safety_settings': _SAFETY_SETTINGS,
+        },
+    },
+]
+_INLINED_TEXT_REQUEST_UNION = {
+    'contents': [{
+        'parts': [{
+            'text': 'high',
+        }],
+        'role': 'user',
+    }],
+    'config': {
+        'response_modalities': ['TEXT'],
+        'system_instruction': 'I say high, you say low',
+        'thinking_config': {
+            'include_thoughts': True,
+            'thinking_budget': 4000,
+        },
+    },
+}
+_INLINED_TEXT_REQUEST = copy.deepcopy(_INLINED_TEXT_REQUEST_UNION)
+_INLINED_TEXT_REQUEST['config']['system_instruction'] = t.t_content(
+    'I say high, you say low'
+)
+_INLINED_TEXT_REQUEST['config']['tools'] = [{'google_search': {}}]
+_INLINED_TEXT_REQUEST['config']['tool_config'] = {
+    'retrieval_config': {'lat_lng': {'latitude': 37.422, 'longitude': -122.084}}
+}
+
+_INLINED_IMAGE_REQUEST = {
+    'contents': [{
+        'parts': [
+            {'text': 'What is in this image?'},
+            {
+                'file_data': {
+                    'file_uri': (
+                        'https://generativelanguage.googleapis.com/v1beta/files/kje1wewvo85z'
+                    ),
+                    'mime_type': 'image/jpeg',
+                },
+            },
+        ],
+        'role': 'user',
+    }],
+    'config': {
+        'temperature': 0.7,
+        'top_p': 0.9,
+        'top_k': 10,
+    },
+}
+_INLINED_VIDEO_REQUEST = {
+    'contents': [{
+        'parts': [
+            {
+                'text': 'Summerize this video.',
+            },
+            {
+                'file_data': {
+                    'file_uri': (
+                        'https://generativelanguage.googleapis.com/v1beta/files/tyvaih24jwje'
+                    ),
+                    'mime_type': 'video/mp4',
+                },
+                'video_metadata': {
+                    'start_offset': '0s',
+                    'end_offset': '5s',
+                    'fps': 3,
+                },
+            },
+        ],
+        'role': 'user',
+    }],
+}
+_IMAGE_PNG_FILE_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/google.png')
+)
+with open(_IMAGE_PNG_FILE_PATH, 'rb') as image_file:
+  image_bytes = image_file.read()
+  image_string = base64.b64encode(image_bytes).decode('utf-8')
+_INLINED_IMAGE_BLOB_REQUEST = types.InlinedRequest(
+    contents=[
+        types.Content(
+            parts=[
+                types.Part(text='What is this image about?'),
+                types.Part(
+                    inline_data=types.Blob(
+                        data=image_string,
+                        mime_type='image/png',
+                    ),
+                ),
+            ],
+            role='user',
+        )
+    ]
+)
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_union_with_inlined_request',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src=_INLINED_REQUESTS,
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported in',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_inlined_request',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': _INLINED_REQUESTS},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_inlined_request_config',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': _INLINED_REQUESTS},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_union_with_inlined_request_system_instruction',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': [_INLINED_TEXT_REQUEST_UNION]},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        has_union=True,
+        exception_if_vertex='supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_image_file',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': [_INLINED_IMAGE_REQUEST]},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_image_blob',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': [_INLINED_IMAGE_BLOB_REQUEST]},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_video_file',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': [_INLINED_VIDEO_REQUEST]},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='supported in',
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create',
+        test_table=test_table,
+    ),
+]
+
+
+@pytest.mark.asyncio
+async def test_async_create(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    batch_job = await client.aio.batches.create(
+        model=_GEMINI_MODEL,
+        src=_INLINED_REQUESTS,
+    )
+    assert batch_job.name.startswith('batches/')
+    assert (
+        batch_job.model == 'models/' + _GEMINI_MODEL
+    )  # Converted to Gemini full name.
@@ -0,0 +1,130 @@
+# 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 batches.create() with Vertex dataset source."""
+
+import re
+
+import pytest
+
+from .. import pytest_helper
+from ... import types
+
+
+_GEMINI_MODEL = 'gemini-2.5-flash'
+_GEMINI_MODEL_FULL_NAME = 'publishers/google/models/gemini-2.5-flash'
+_OUTPUT_VERTEX_DATASET_DISPLAY_NAME = 'test_batch_output'
+_VERTEX_DATASET_INPUT_NAME = (
+    'projects/vertex-sdk-dev/locations/us-central1/datasets/7857316250517504000'
+)
+_DISPLAY_NAME = 'test_batch'
+
+_BQ_OUTPUT_PREFIX = (
+    'bq://vertex-sdk-dev.unified_genai_tests_batches.generate_content_output'
+)
+_VERTEX_DATASET_DESTINATION = types.VertexMultimodalDatasetDestination(
+    bigquery_destination=_BQ_OUTPUT_PREFIX,
+    display_name=_OUTPUT_VERTEX_DATASET_DISPLAY_NAME,
+)
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_union_generate_content_with_vertex_dataset_name',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL_FULL_NAME,
+            src=_VERTEX_DATASET_INPUT_NAME,
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': {
+                    'vertex_dataset': _VERTEX_DATASET_DESTINATION,
+                    'format': 'vertex-dataset',
+                },
+            },
+        ),
+        exception_if_mldev='not supported in Gemini API',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_generate_content_with_vertex_dataset_source',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL_FULL_NAME,
+            src=types.BatchJobSource(
+                vertex_dataset_name=_VERTEX_DATASET_INPUT_NAME,
+                format='vertex-dataset',
+            ),
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': {
+                    'vertex_dataset': _VERTEX_DATASET_DESTINATION,
+                    'format': 'vertex-dataset',
+                },
+            },
+        ),
+        exception_if_mldev='one of',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_generate_content_with_vertex_dataset_source_dict',
+        parameters=types._CreateBatchJobParameters(
+            model=_GEMINI_MODEL_FULL_NAME,
+            src={
+                'vertex_dataset_name': _VERTEX_DATASET_INPUT_NAME,
+                'format': 'vertex-dataset',
+            },
+            config={
+                'display_name': _DISPLAY_NAME,
+                'dest': {
+                    'vertex_dataset': _VERTEX_DATASET_DESTINATION,
+                    'format': 'vertex-dataset',
+                },
+            },
+        ),
+        exception_if_mldev='one of',
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create',
+        test_table=test_table,
+    ),
+]
+
+
+@pytest.mark.asyncio
+async def test_async_create(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    batch_job = await client.aio.batches.create(
+        model=_GEMINI_MODEL,
+        src=_VERTEX_DATASET_INPUT_NAME,
+        config={
+            'dest': {
+                'vertex_dataset': _VERTEX_DATASET_DESTINATION,
+                'format': 'vertex-dataset',
+            },
+        },
+    )
+
+    assert batch_job.name.startswith('projects/')
+    assert (
+        batch_job.model == _GEMINI_MODEL_FULL_NAME
+    )  # Converted to Vertex full name.
+    assert batch_job.src.vertex_dataset_name == _VERTEX_DATASET_INPUT_NAME
+    assert batch_job.src.format == 'vertex-dataset'
@@ -0,0 +1,86 @@
+# 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 batches.delete()."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+_BATCH_JOB_NAME = '7085929781874655232'
+_BATCH_JOB_FULL_RESOURCE_NAME = (
+    'projects/964831358985/locations/us-central1/'
+    f'batchPredictionJobs/{_BATCH_JOB_NAME}'
+)
+_INVALID_BATCH_JOB_NAME = 'invalid_name'
+_MLDEV_BATCH_OPERATION_NAME = 'batches/70h2jo0ic2t1zejyl0p4jgi8mk1gj0wvjusv'
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_delete_batch_job',
+        parameters=types._DeleteBatchJobParameters(
+            name=_BATCH_JOB_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+        skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_delete_batch_job_full_resource_name',
+        override_replay_id='test_delete_batch_job',
+        parameters=types._DeleteBatchJobParameters(
+            name=_BATCH_JOB_FULL_RESOURCE_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+        skip_in_api_mode=('Cannot cancel a batch job multiple times in Vertex'),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_delete_batch_job_with_invalid_name',
+        parameters=types._DeleteBatchJobParameters(
+            name=_INVALID_BATCH_JOB_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+        exception_if_vertex='Invalid batch job name',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_delete_batch_operation',
+        parameters=types._DeleteBatchJobParameters(
+            name=_MLDEV_BATCH_OPERATION_NAME,
+        ),
+        exception_if_vertex='Invalid batch job name',
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='batches.delete',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async_delete(client):
+  if client.vertexai:
+    name = _BATCH_JOB_NAME
+  else:
+    name = _MLDEV_BATCH_OPERATION_NAME
+
+  delete_job = await client.aio.batches.delete(name=name)
+  assert delete_job
@@ -0,0 +1,159 @@
+# 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 batches._create_embeddings()"""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+_MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME = (
+    'batches/wdx71o8cgbzoa6gg3be1mg7g8ulrhapcjgo3'
+)
+_MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME = (
+    'batches/507oatd242het8ox60pwsmn7tcmtkrj8itff'
+)
+
+_DISPLAY_NAME = 'test_batch'
+_MLDEV_EMBEDDING_MODEL = 'gemini-embedding-001'
+_EMBED_CONTENT_FILE_NAME = 'files/mq9e3mg3u2y5'
+_INLINED_EMBED_CONTENT_REQUESTS = {
+    'config': {'output_dimensionality': 64},
+    'contents': [
+        {
+            'parts': [{
+                'text': '1',
+            }],
+        },
+        {
+            'parts': [{
+                'text': '2',
+            }],
+        },
+        {
+            'parts': [{
+                'text': '3',
+            }],
+        },
+    ],
+}
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_from_inlined',
+        parameters=types._CreateEmbeddingsBatchJobParameters(
+            model=_MLDEV_EMBEDDING_MODEL,
+            src={'inlined_requests': _INLINED_EMBED_CONTENT_REQUESTS},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex=(
+            'Gemini Enterprise Agent Platform'
+        ),
+    ),
+]
+
+pytestmark = [
+    pytest.mark.usefixtures('mock_timestamped_unique_name'),
+    pytest_helper.setup(
+        file=__file__,
+        globals_for_file=globals(),
+        test_method='batches.create_embeddings',
+        test_table=test_table,
+        http_options={
+            'api_version': 'v1alpha',
+            'base_url': (
+                'https://autopush-generativelanguage.sandbox.googleapis.com'
+            ),
+        },
+    ),
+]
+
+
+@pytest.mark.asyncio
+async def test_async_from_inline(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    batch_job = await client.aio.batches.create_embeddings(
+        model=_MLDEV_EMBEDDING_MODEL,
+        src={'inlined_requests': _INLINED_EMBED_CONTENT_REQUESTS},
+    )
+    assert batch_job.name.startswith('batches/')
+
+
+def test_from_file(client):
+  """Tests creating a batch job with an embedding file name."""
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    batch_job = client.batches.create_embeddings(
+        model=_MLDEV_EMBEDDING_MODEL,
+        src={'file_name': _EMBED_CONTENT_FILE_NAME},
+        config={
+            'display_name': _DISPLAY_NAME,
+        },
+    )
+    assert batch_job.name.startswith('batches/')
+    assert batch_job.model == 'models/' + _MLDEV_EMBEDDING_MODEL
+
+
+@pytest.mark.asyncio
+async def test_async_from_file(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    batch_job = await client.aio.batches.create_embeddings(
+        model=_MLDEV_EMBEDDING_MODEL,
+        src={'file_name': _EMBED_CONTENT_FILE_NAME},
+        config={
+            'display_name': _DISPLAY_NAME,
+        },
+    )
+    assert batch_job.name.startswith('batches/')
+    assert (
+        batch_job.model == 'models/' + _MLDEV_EMBEDDING_MODEL
+    )  # Converted to Gemini full name.
+
+
+def test_get_inline(client):
+  """Tests getting a batch job that used inline requests."""
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    name = _MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME
+    batch_job = client.batches.get(name=name)
+    assert batch_job.dest.inlined_embed_content_responses is not None
+
+
+@pytest.mark.asyncio
+async def test_async_get_inline(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    name = _MLDEV_EMBEDDING_BATCH_INLINE_OPERATION_NAME
+    batch_job = await client.aio.batches.get(name=name)
+
+    assert batch_job.dest.inlined_embed_content_responses is not None
+
+
+def test_get_file(client):
+  """Tests getting a batch job that used a file source."""
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    name = _MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME
+    batch_job = client.batches.get(name=name)
+    assert batch_job.dest.file_name is not None
+
+
+@pytest.mark.asyncio
+async def test_async_get_file(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    name = _MLDEV_EMBEDDING_BATCH_FILE_OPERATION_NAME
+    batch_job = await client.aio.batches.get(name=name)
+
+    assert batch_job.dest.file_name is not None
@@ -0,0 +1,91 @@
+# 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 batches.get()."""
+
+import re
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+
+# Vertex AI batch job name.
+_BATCH_JOB_NAME = '5798522612028014592'
+_BATCH_JOB_FULL_RESOURCE_NAME = (
+    'projects/964831358985/locations/us-central1/'
+    f'batchPredictionJobs/{_BATCH_JOB_NAME}'
+)
+# MLDev batch operation name.
+_MLDEV_BATCH_OPERATION_NAME = 'batches/z2p8ksus4lyxt25rntl3fpd67p2niw4hfij5'
+_INVALID_BATCH_JOB_NAME = 'invalid_name'
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_get_batch_job',
+        parameters=types._GetBatchJobParameters(
+            name=_BATCH_JOB_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_batch_operation',
+        parameters=types._GetBatchJobParameters(
+            name=_MLDEV_BATCH_OPERATION_NAME,
+        ),
+        exception_if_vertex='Invalid batch job name',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_batch_job_with_invalid_name',
+        parameters=types._GetBatchJobParameters(
+            name=_INVALID_BATCH_JOB_NAME,
+        ),
+        exception_if_mldev='Invalid batch job name',
+        exception_if_vertex='Invalid batch job name',
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='batches.get',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async_get(client):
+  if client.vertexai:
+    name = _BATCH_JOB_NAME
+  else:
+    name = _MLDEV_BATCH_OPERATION_NAME
+  batch_job = await client.aio.batches.get(name=name)
+
+  assert batch_job
+
+
+@pytest.mark.asyncio
+async def test_async_get_with_multimodal_dataset_output(client):
+  if client.vertexai:
+    name = _BATCH_JOB_NAME
+    batch_job = await client.aio.batches.get(name=name)
+
+    assert re.match(
+        r'^projects/[^/]+/locations/[^/]+/datasets/[^/]+$',
+        batch_job.output_info.vertex_multimodal_dataset_name,
+    )
@@ -0,0 +1,79 @@
+# 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 batches.list()."""
+
+import pytest
+
+from ... import errors
+from ... import types
+from .. import pytest_helper
+
+
+# All tests will be run for both Vertex and MLDev.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_list_batch_jobs',
+        parameters=types._ListBatchJobsParameters(
+            config=types.ListBatchJobsConfig(
+                page_size=5,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_list_batch_jobs_with_config',
+        parameters=types._ListBatchJobsParameters(
+            config=types.ListBatchJobsConfig(
+                filter='display_name:"genai_*"',
+                page_size=5,
+            ),
+        ),
+        exception_if_mldev='filter parameter is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode.',
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='batches.list',
+    test_table=test_table,
+)
+
+
+def test_pager(client):
+  batch_jobs = client.batches.list(config={'page_size': 10})
+  assert 'content-type' in batch_jobs.sdk_http_response.headers
+  assert batch_jobs.name == 'batch_jobs'
+  assert batch_jobs.page_size == 10
+  assert len(batch_jobs) <= 10
+
+  # Iterate through the first page. Otherwise, too many batches are returned.
+  for _ in batch_jobs.page:
+    pass
+
+
+@pytest.mark.asyncio
+async def test_async_pager(client):
+  batch_jobs = await client.aio.batches.list(config={'page_size': 10})
+
+  assert 'Content-Type' in batch_jobs.sdk_http_response.headers
+  assert batch_jobs.name == 'batch_jobs'
+  assert batch_jobs.page_size == 10
+  assert len(batch_jobs) <= 10
+
+  # Iterate through the first page. Otherwise, too many batches are returned.
+  for _ in batch_jobs.page:
+    pass