google-genai 1.56.0__py3-none-any.whl → 1.58.0__py3-none-any.whl

google/genai/tests/models/test_update.py

@@ -0,0 +1,95 @@
+# 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.
+#
+
+
+import pytest
+from ... import errors
+from ... import types
+from .. import pytest_helper
+
+test_http_options = {'headers': {'test': 'headers'}}
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_vertex_tuned_models_update',
+        parameters=types._UpdateModelParameters(
+            model='models/2171259487439028224',
+            config={
+                'description': (
+                    'My SupervisedTuningJob'
+                ),
+                'default_checkpoint_id': '8',
+            },
+        ),
+        exception_if_mldev='404',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_vertex_tuned_models_update_with_http_options_in_method',
+        parameters=types._UpdateModelParameters(
+            model='models/2171259487439028224',
+            config={
+                'description': (
+                    'My SupervisedTuningJob'
+                ),
+                'default_checkpoint_id': '8',
+                'http_options': test_http_options,
+            },
+        ),
+        exception_if_mldev='404',
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.update',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async_update_tuned_model(client):
+  if client._api_client.vertexai:
+    with pytest.raises(errors.ClientError) as e:
+      await client.aio.models.update(
+          model='tunedModels/generatenum5443-ekrw7ie9wis23zbeogbw6jq8',
+          config={
+              'description': 'My tuned gemini model',
+              'http_options': test_http_options,
+          },
+      )
+    assert '404' in str(e)
+
+
+@pytest.mark.asyncio
+async def test_async_update_model(client):
+  if client._api_client.vertexai:
+    response = await client.aio.models.update(
+        model='models/2171259487439028224',
+        config={
+            'display_name': 'My tuned gemini model',
+            'http_options': test_http_options,
+        },
+    )
+  else:
+    with pytest.raises(errors.ClientError) as e:
+      await client.aio.models.update(
+          model='models/2171259487439028224',
+          config={
+              'display_name': 'My tuned gemini model',
+              'http_options': test_http_options,
+          },
+      )
+    assert '404' in str(e)

google/genai/tests/models/test_upscale_image.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 upscale_image."""
+
+import os
+
+from pydantic import ValidationError
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+IMAGEN_MODEL_LATEST = 'imagen-4.0-upscale-preview'
+
+IMAGE_FILE_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/bridge1.png')
+)
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_upscale_no_config',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types.UpscaleImageParameters(
+            model=IMAGEN_MODEL_LATEST,
+            image=types.Image.from_file(location=IMAGE_FILE_PATH),
+            upscale_factor='x2',
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_upscale',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types.UpscaleImageParameters(
+            model=IMAGEN_MODEL_LATEST,
+            image=types.Image.from_file(location=IMAGE_FILE_PATH),
+            upscale_factor='x2',
+            config=types.UpscaleImageConfig(
+                include_rai_reason=True,
+                person_generation=types.PersonGeneration.ALLOW_ADULT,
+                safety_filter_level=types.SafetyFilterLevel.BLOCK_LOW_AND_ABOVE,
+                output_mime_type='image/jpeg',
+                output_compression_quality=80,
+                enhance_input_image=True,
+                image_preservation_factor=0.6,
+                labels={'imagen_label_key': 'upscale_image'}
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_upscale_gcs',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types.UpscaleImageParameters(
+            model=IMAGEN_MODEL_LATEST,
+            image=types.Image.from_file(location=IMAGE_FILE_PATH),
+            upscale_factor='x2',
+            config=types.UpscaleImageConfig(
+                output_gcs_uri='gs://genai-sdk-tests/temp/images/',
+            ),
+        ),
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.upscale_image',
+    test_table=test_table,
+)
+
+
+def test_upscale_extra_config_parameters(client):
+  # MLDev currently does not support upscale_image, but the ValidationError
+  # occurs before the ValueError.
+  try:
+    # User is not allowed to set mode or number_of_images
+    client.models.upscale_image(
+        model=IMAGEN_MODEL_LATEST,
+        image=types.Image.from_file(location=IMAGE_FILE_PATH),
+        upscale_factor='x2',
+        config={
+            'mode': 'upscale',
+            'number_of_images': 1,
+        }
+    )
+    # Should never reach this.
+    assert False
+  except Exception as e:
+    assert isinstance(e, ValidationError)
+    assert 'Extra inputs are not permitted' in str(e)
+
+
+@pytest.mark.asyncio
+async def test_upscale_async(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    response = await client.aio.models.upscale_image(
+        model=IMAGEN_MODEL_LATEST,
+        image=types.Image.from_file(location=IMAGE_FILE_PATH),
+        upscale_factor='x2',
+        config=types.UpscaleImageConfig(
+            person_generation=types.PersonGeneration.ALLOW_ADULT,
+            safety_filter_level=types.SafetyFilterLevel.BLOCK_LOW_AND_ABOVE,
+            include_rai_reason=True,
+            output_mime_type='image/jpeg',
+            output_compression_quality=80,
+            enhance_input_image=True,
+            image_preservation_factor=0.6,
+        ),
+    )
+    assert response.generated_images[0].image.image_bytes
+
+
+@pytest.mark.asyncio
+async def test_upscale_gcs_async(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    response = await client.aio.models.upscale_image(
+        model=IMAGEN_MODEL_LATEST,
+        image=types.Image.from_file(location=IMAGE_FILE_PATH),
+        upscale_factor='x2',
+        config=types.UpscaleImageConfig(
+            output_gcs_uri='gs://genai-sdk-tests/temp/images/',
+        ),
+    )
+    assert response.generated_images[0].image.gcs_uri
+
+
+@pytest.mark.asyncio
+async def test_upscale_extra_config_parameters_async(client):
+  # MLDev currently does not support upscale_image, but the ValidationError
+  # occurs before the ValueError.
+  try:
+    # User is not allowed to set mode or number_of_images
+    await client.aio.models.upscale_image(
+        model=IMAGEN_MODEL_LATEST,
+        image=types.Image.from_file(location=IMAGE_FILE_PATH),
+        upscale_factor='x2',
+        config={
+            'mode': 'upscale',
+            'number_of_images': 1,
+        },
+    )
+    # Should never reach this.
+    assert False
+  except Exception as e:
+    assert isinstance(e, ValidationError)
+    assert 'Extra inputs are not permitted' in str(e)

google/genai/tests/operations/__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 operations.get."""

google/genai/tests/operations/test_get.py

@@ -0,0 +1,38 @@
+ # 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 operations._get."""
+
+from .. import pytest_helper
+from ... import types
+
+
+def test_project_operation_get(client):
+  test_operation_id = '3787416390563004416'
+  if client._api_client.vertexai:
+    operation = client.operations._get(
+        operation_id=test_operation_id
+    )
+    assert operation.name.endswith(test_operation_id)
+    assert operation.done
+    assert isinstance(operation, types.ProjectOperation)
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='operations._get',
+)

google/genai/tests/public_samples/__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.
+#
+
+
+"""Public Samples and Tests for the Google GenAI SDK."""

google/genai/tests/public_samples/test_gemini_text_only.py

@@ -0,0 +1,34 @@
+# 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.
+#
+
+
+"""Test for the code sample for Gemini text-only request."""
+
+from .. import pytest_helper
+
+pytestmark = pytest_helper.setup(file=__file__)
+
+
+def test_sample(client):
+  # [START generativeaionvertexai_gemini_text_only]
+  response = client.models.generate_content(
+      model="gemini-2.5-flash",
+      contents=(
+          "What's a good name for a flower shop that specializes in selling"
+          " bouquets of dried flowers?"
+      ),
+  )
+  print(response.text)
+  # [END generativeaionvertexai_gemini_text_only]

google/genai/tests/pytest_helper.py

@@ -0,0 +1,246 @@
+# 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.
+#
+
+
+import contextlib
+import json
+import os
+import pathlib
+from typing import Any, Optional
+from pydantic import BaseModel, Field, SerializeAsAny
+import pytest
+import re
+from .. import _common
+from .. import _replay_api_client
+from .. import types
+from .._api_client import HttpOptions
+
+is_api_mode = "config.getoption('--mode') == 'api'"
+
+
+class TestTableItem(types.TestTableItem):
+  # This is not a test suite class.
+  __test__ = False
+
+  # Overrides to support Pydantic models.
+  parameters: SerializeAsAny[BaseModel] = Field(
+      description="""The parameters to the test. Use pydantic models.""",
+  )
+
+
+def base_test_function(
+    client,
+    use_vertex: bool,
+    replays_prefix: str,
+    test_method: str,
+    test_table_item: TestTableItem,
+    globals_for_file: dict[str, Any],
+):
+  replay_id = (
+      test_table_item.override_replay_id
+      if test_table_item.override_replay_id
+      else test_table_item.name
+  )
+  api_type = 'vertex' if use_vertex else 'mldev'
+  replay_id = f'{replays_prefix}/{replay_id}.{api_type}'
+  client._api_client.initialize_replay_session(replay_id)
+  # vars().copy() provides a shallow copy of the parameters.
+  parameters_dict = vars(test_table_item.parameters).copy()
+  try:
+    if '.' in test_method:
+      method_name_parts = test_method.split('.')
+      module_path_parts = method_name_parts[:-1]
+      method_name = method_name_parts[-1]
+      # Iterates the nested module starting from client
+      current_object = client
+      for part in module_path_parts:
+        current_object = getattr(current_object, part)
+      method = getattr(current_object, method_name)
+      method(**parameters_dict)
+    else:
+      custom_method = globals_for_file[test_method]
+      custom_method(client, test_table_item.parameters)
+    # Should not reach here if expecting an exception.
+    if test_table_item.exception_if_mldev and not client._api_client.vertexai:
+      assert False, 'Should have raised exception in MLDev.'
+    elif test_table_item.exception_if_vertex and client._api_client.vertexai:
+      assert False, 'Should have raised exception in Vertex.'
+    client._api_client.close()
+  except Exception as e:
+    if test_table_item.exception_if_mldev and not client._api_client.vertexai:
+      if test_table_item.exception_if_mldev not in str(e):
+        raise AssertionError(
+            f"'{test_table_item.exception_if_mldev}' not in '{str(e)}'"
+        ) from e
+    elif test_table_item.exception_if_vertex and client._api_client.vertexai:
+      if test_table_item.exception_if_vertex not in str(e):
+        raise AssertionError(
+            f"'{test_table_item.exception_if_vertex}' not in '{str(e)}'"
+        ) from e
+    else:
+      raise e
+
+
+def create_test_for_table_item(
+    globals_for_file: dict[str, Any],
+    test_method: str,
+    test_table_item: TestTableItem,
+):
+  return lambda client, use_vertex, replays_prefix: base_test_function(
+      client,
+      use_vertex,
+      replays_prefix,
+      test_method,
+      test_table_item,
+      globals_for_file,
+  )
+
+
+def create_test_for_table(
+    globals_for_file: dict[str, Any],
+    test_method: str,
+    test_table: list[TestTableItem],
+):
+  for test_table_item in test_table:
+    if test_table_item.has_union:
+      assert test_table_item.name.startswith('test_union_'), f"""
+      test: {test_table_item.name}
+      When has_union is true, it must be a test that explicitly test the
+      non-canonical type in the union.
+      For all other tests, use transformers to convert data to the canonical
+      type of the field so it can be tested in other languages.
+      If this test is truly about testing the non-canonical type of the union,
+      rename the test to start with 'test_union_'. E.g.
+      test_union_contents_is_string()"""
+    globals_for_file[test_table_item.name] = create_test_for_table_item(
+        globals_for_file, test_method, test_table_item
+    )
+
+
+# Sets up the test framework.
+# file: Always use __file__
+# globals_for_file: Always use globals()
+# test_table: The test table for the file.
+#     Use test table over individual tests whenever possible.
+#     Tests built with test_table will run in all other languages automatically.
+#     Otherwise, you will need to write tests in all other languages manually.
+def setup(
+    *,
+    file: str,
+    globals_for_file: Optional[dict[str, Any]] = None,
+    test_method: Optional[str] = None,
+    test_table: Optional[list[TestTableItem]] = None,
+    http_options: Optional[HttpOptions] = None,
+):
+  """Generates parameterization for tests, run for both Vertex and MLDev."""
+  replays_directory = (
+      file.replace(os.path.dirname(__file__), 'tests')
+      .replace('.py', '')
+      .replace('/test_', '/')
+  )
+
+  is_tap_mode = False
+  if os.environ.get('UNITTEST_ON_FORGE', None) is not None:
+    is_tap_mode = True
+
+  if test_table:
+    create_test_for_table(globals_for_file, test_method, test_table)
+
+  if test_table and not is_tap_mode:
+    replays_root_directory = os.environ.get(
+        'GOOGLE_GENAI_REPLAYS_DIRECTORY', None
+    )
+    if replays_root_directory is None:
+      raise ValueError(
+          'GOOGLE_GENAI_REPLAYS_DIRECTORY environment variable is not set'
+      )
+    abs_replay_directory = os.path.join(
+        replays_root_directory, replays_directory
+    )
+    test_table_file_path = os.path.join(
+        abs_replay_directory, '_test_table.json'
+    )
+
+    pathlib.Path(abs_replay_directory).mkdir(parents=True, exist_ok=True)
+    assert isinstance(
+        test_table[0].parameters, BaseModel
+    ), f'{test_table_file_path} parameters must be a BaseModel.'
+    test_table_file = types.TestTableFile(
+        comment='Auto-generated. Do not edit.',
+        test_method=test_method,
+        parameter_names=list(test_table[0].parameters.model_fields.keys()),
+        test_table=test_table,
+    )
+    os.makedirs(os.path.dirname(test_table_file_path), exist_ok=True)
+
+    with open(test_table_file_path, 'w') as f:
+      f.write(
+          test_table_file.model_dump_json(exclude_none=True, by_alias=True, indent=2),
+      )
+
+  # Add fixture for requested client option.
+  return pytest.mark.parametrize(
+      'use_vertex, replays_prefix, http_options',
+      [
+          (True, replays_directory, http_options),
+          (False, replays_directory, http_options),
+      ],
+  )
+
+
+def exception_if_mldev(client, exception_type: type[Exception]):
+  if client._api_client.vertexai:
+    return contextlib.nullcontext()
+  else:
+    return pytest.raises(exception_type)
+
+
+def exception_if_vertex(client, exception_type: type[Exception]):
+  if client._api_client.vertexai:
+    return pytest.raises(exception_type)
+  else:
+    return contextlib.nullcontext()
+
+
+def snake_to_camel(snake_str: str) -> str:
+  """Converts a snake_case string to CamelCase."""
+  return re.sub(r'_([a-zA-Z])', lambda match: match.group(1).upper(), snake_str)
+
+
+def camel_to_snake(camel_str: str) -> str:
+  """Converts a CamelCase string to snake_case."""
+  return re.sub(r'([A-Z])', r'_\1', camel_str).lower().lstrip('_')
+
+
+def get_value_ignore_key_case(obj, key):
+  """Returns the value of the key in the object, converting to camelCase or snake_case if necessary."""
+  return obj.get(snake_to_camel(key), obj.get(camel_to_snake(key), None))
+
+
+def camel_to_snake_all_keys(data):
+  """Converts all keys in a dictionary or list to snake_case."""
+  if isinstance(data, dict):
+    new_dict = {}
+    for key, value in data.items():
+      if isinstance(key, str):
+        new_key = camel_to_snake(key)
+      else:
+        new_key = key
+      new_dict[new_key] = camel_to_snake_all_keys(value)
+    return new_dict
+  elif isinstance(data, list):
+    return [camel_to_snake_all_keys(item) for item in data]
+  else:
+    return data

google/genai/tests/shared/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+#
+
+

google/genai/tests/shared/batches/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+#