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

google/genai/tests/batches/test_create_with_inlined_requests.py

@@ -0,0 +1,255 @@
+# 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_REQUEST = {
+    'contents': [{
+        'parts': [{
+            'text': 'Hello!',
+        }],
+        'role': 'user',
+    }],
+    'metadata': {
+        'key': 'request-1',
+    },
+    '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_REQUEST],
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='not supported in Vertex',
+        has_union=True,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_inlined_request',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': [_INLINED_REQUEST]},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='not supported',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_inlined_request_config',
+        parameters=types._CreateBatchJobParameters(
+            model=_MLDEV_GEMINI_MODEL,
+            src={'inlined_requests': [_INLINED_TEXT_REQUEST]},
+            config={
+                'display_name': _DISPLAY_NAME,
+            },
+        ),
+        exception_if_vertex='not supported',
+    ),
+    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='not supported',
+    ),
+    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='not supported',
+    ),
+    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='not supported',
+    ),
+    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='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=_GEMINI_MODEL,
+        src=[_INLINED_REQUEST],
+    )
+    assert batch_job.name.startswith('batches/')
+    assert (
+        batch_job.model == 'models/' + _GEMINI_MODEL
+    )  # Converted to Gemini full name.

google/genai/tests/batches/test_delete.py

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

google/genai/tests/batches/test_embedding.py

@@ -0,0 +1,157 @@
+# 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='Vertex AI does not support',
+    ),
+]
+
+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

google/genai/tests/batches/test_get.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.get()."""
+
+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/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_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

google/genai/tests/batches/test_list.py

@@ -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 not supported in Gemini API',
+    ),
+]
+
+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

google/genai/tests/caches/__init__.py

@@ -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 caches module."""

google/genai/tests/caches/constants.py

@@ -0,0 +1,29 @@
+# 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.
+#
+
+
+"""Constants for caches tests."""
+
+CACHED_CONTENT_NAME_MLDEV = 'cachedContents/o239k1gxzz0juy9wqstndhncr85krehehf551hqh'
+CACHED_CONTENT_NAME_VERTEX = 'cachedContents/2164089915711684608'
+
+VERTEX_HTTP_OPTIONS = {
+    'api_version': 'v1beta1',
+    'base_url': 'https://us-central1-aiplatform.googleapis.com/',
+}
+MLDEV_HTTP_OPTIONS = {
+    'api_version': 'v1beta',
+    'base_url': 'https://generativelanguage.googleapis.com/',
+}