google-genai 1.53.0__py3-none-any.whl → 1.55.0__py3-none-any.whl

google/genai/tests/batches/test_cancel.py

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

google/genai/tests/batches/test_create.py

@@ -0,0 +1,78 @@
+# 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='not supported in Gemini API',
+        exception_if_vertex='Unsupported destination',
+        has_union=True,
+    ),
+]
+
+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,
+    ),
+]

google/genai/tests/batches/test_create_with_bigquery.py

@@ -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 = 'text-embedding-004'
+_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/text-embedding-004'
+_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=_EMBEDDING_MODEL_FULL_NAME,
+            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'

google/genai/tests/batches/test_create_with_file.py

@@ -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='not 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='not 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.

google/genai/tests/batches/test_create_with_gcs.py

@@ -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 = 'text-embedding-004'
+_EMBEDDING_MODEL_FULL_NAME = 'publishers/google/models/text-embedding-004'
+_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=_EMBEDDING_MODEL,
+            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'