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

google/genai/tests/models/test_compute_tokens.py

@@ -0,0 +1,120 @@
+# 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 _transformers as t
+from ... import types
+from .. import pytest_helper
+from . import constants
+
+_COMPUTE_TOKENS_PARAMS = types._ComputeTokensParameters(
+    model='gemini-2.5-flash',
+    contents=[t.t_content('Tell me a story in 300 words.')],
+)
+_COMPUTE_TOKENS_PARAMS_VERTEX_CUSTOM_URL = types._ComputeTokensParameters(
+    model='gemini-2.5-flash',
+    contents=[t.t_content('Tell me a story in 300 words.')],
+    config={'http_options': constants.VERTEX_HTTP_OPTIONS},
+)
+_COMPUTE_TOKENS_PARAMS_MLDEV_CUSTOM_URL = types._ComputeTokensParameters(
+    model='gemini-2.5-flash',
+    contents=[t.t_content('Tell me a story in 300 words.')],
+    config={'http_options': constants.MLDEV_HTTP_OPTIONS},
+)
+_UNICODE_STRING = '这是一条unicode测试🤪❤★'
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_compute_tokens',
+        exception_if_mldev='only supported in',
+        parameters=types._ComputeTokensParameters(
+            model='gemini-2.5-flash',
+            contents=[t.t_content('Tell me a story in 300 words.')],
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_compute_tokens_vertex_custom_url',
+        parameters=_COMPUTE_TOKENS_PARAMS_VERTEX_CUSTOM_URL,
+        exception_if_mldev='only supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_compute_tokens_mldev_custom_url',
+        parameters=_COMPUTE_TOKENS_PARAMS_MLDEV_CUSTOM_URL,
+        exception_if_vertex='404',
+        exception_if_mldev='only supported in',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_compute_tokens_unicode',
+        exception_if_mldev='only supported in',
+        parameters=types._ComputeTokensParameters(
+            model='gemini-2.5-flash', contents=[t.t_content(_UNICODE_STRING)]
+        ),
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.compute_tokens',
+    test_table=test_table,
+)
+
+
+def test_token_bytes_deserialization(client):
+  if client._api_client.vertexai:
+    response = client.models.compute_tokens(
+        model=_COMPUTE_TOKENS_PARAMS.model,
+        contents=_UNICODE_STRING,
+    )
+    decoded_tokens = b''.join(response.tokens_info[0].tokens)
+    assert (
+        decoded_tokens
+        == b'\xe8\xbf\x99\xe6\x98\xaf\xe4\xb8\x80\xe6\x9d\xa1unicode\xe6\xb5\x8b\xe8\xaf\x95\xf0\x9f\xa4\xaa\xe2\x9d\xa4\xe2\x98\x85'
+    )
+    assert decoded_tokens.decode('utf-8') == _UNICODE_STRING
+
+
+@pytest.mark.asyncio
+async def test_async(client):
+  if client._api_client.vertexai:
+    response = await client.aio.models.compute_tokens(
+        model=_COMPUTE_TOKENS_PARAMS.model,
+        contents=_COMPUTE_TOKENS_PARAMS.contents,
+    )
+    assert response
+  else:
+    with pytest.raises(Exception):
+      await client.aio.models.compute_tokens(
+          model=_COMPUTE_TOKENS_PARAMS.model,
+          contents=_COMPUTE_TOKENS_PARAMS.contents,
+      )
+
+
+def test_different_model_names(client):
+  if client._api_client.vertexai:
+    response1 = client.models.compute_tokens(
+        model='gemini-2.5-flash', contents=_COMPUTE_TOKENS_PARAMS.contents
+    )
+    assert response1
+    response3 = client.models.compute_tokens(
+        model='publishers/google/models/gemini-2.5-flash',
+        contents=_COMPUTE_TOKENS_PARAMS.contents,
+    )
+    assert response3
+    response4 = client.models.compute_tokens(
+        model='projects/vertexsdk/locations/us-central1/publishers/google/models/gemini-2.5-flash',
+        contents=_COMPUTE_TOKENS_PARAMS.contents,
+    )
+    assert response4

google/genai/tests/models/test_count_tokens.py

@@ -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.
+#
+
+
+import copy
+import pytest
+from ... import _transformers as t
+from ... import errors
+from ... import types
+from .. import pytest_helper
+from . import constants
+
+_COUNT_TOKENS_PARAMS = types._CountTokensParameters(
+    model='gemini-2.5-flash',
+    contents=[t.t_content('Tell me a story in 300 words.')],
+)
+
+_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION = copy.deepcopy(
+    _COUNT_TOKENS_PARAMS
+)
+_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION.config = {
+    'system_instruction': t.t_content('you are a chatbot.')
+}
+
+_COUNT_TOKENS_PARAMS_WITH_TOOLS = copy.deepcopy(_COUNT_TOKENS_PARAMS)
+_COUNT_TOKENS_PARAMS_WITH_TOOLS.config = {
+    'tools': [{'google_search_retrieval': {}}]
+}
+
+_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG = copy.deepcopy(
+    _COUNT_TOKENS_PARAMS
+)
+_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG.config = {
+    'generation_config': {'max_output_tokens': 50}
+}
+
+_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL = copy.deepcopy(_COUNT_TOKENS_PARAMS)
+_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL.config = {
+    'http_options': constants.VERTEX_HTTP_OPTIONS
+}
+_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL = copy.deepcopy(_COUNT_TOKENS_PARAMS)
+_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL.config = {
+    'http_options': constants.MLDEV_HTTP_OPTIONS
+}
+
+
+# TODO(b/378952792): MLDev count_tokens needs to merge contents and model
+# param into generateContentRequest field.
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_count_tokens',
+        parameters=_COUNT_TOKENS_PARAMS,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_count_tokens_vertex_custom_url',
+        parameters=_COUNT_TOKENS_PARAMS_VERTEX_CUSTOM_URL,
+        exception_if_mldev='404',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_count_tokens_mldev_custom_url',
+        parameters=_COUNT_TOKENS_PARAMS_MLDEV_CUSTOM_URL,
+        exception_if_vertex='404',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_count_tokens_with_system_instruction',
+        exception_if_mldev='not supported',
+        parameters=_COUNT_TOKENS_PARAMS_WITH_SYSTEM_INSTRUCTION,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_count_tokens_with_tools',
+        exception_if_mldev='not supported',
+        parameters=_COUNT_TOKENS_PARAMS_WITH_TOOLS,
+    ),
+    pytest_helper.TestTableItem(
+        name='test_count_tokens_with_generation_config',
+        exception_if_mldev='not supported',
+        parameters=_COUNT_TOKENS_PARAMS_WITH_GENERATION_CONFIG,
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.count_tokens',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async(client):
+  response = await client.aio.models.count_tokens(
+      model=_COUNT_TOKENS_PARAMS.model, contents=_COUNT_TOKENS_PARAMS.contents
+  )
+  assert response
+
+
+def test_different_model_names(client):
+  if client._api_client.vertexai:
+    response1 = client.models.count_tokens(
+        model='gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
+    )
+    assert response1
+    response3 = client.models.count_tokens(
+        model='publishers/google/models/gemini-2.5-flash',
+        contents=_COUNT_TOKENS_PARAMS.contents,
+    )
+    assert response3
+    response4 = client.models.count_tokens(
+        model='projects/vertexsdk/locations/us-central1/publishers/google/models/gemini-2.5-flash',
+        contents=_COUNT_TOKENS_PARAMS.contents,
+    )
+    assert response4
+  else:
+    response1 = client.models.count_tokens(
+        model='gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
+    )
+    assert response1
+    response2 = client.models.count_tokens(
+        model='models/gemini-2.5-flash', contents=_COUNT_TOKENS_PARAMS.contents
+    )
+    assert response2
+
+
+def test_extra_body(client):
+  config = {
+      'http_options': {
+          'extra_body': {
+              'systemInstruction': {
+                  'parts': [{'text': 'you are a chatbot.'}],
+                  'role': 'user',
+              }
+          }
+      }
+  }
+  if client._api_client.vertexai:
+    response = client.models.count_tokens(
+        model=_COUNT_TOKENS_PARAMS.model,
+        contents=_COUNT_TOKENS_PARAMS.contents,
+        config=config,
+    )
+    assert response.total_tokens
+  else:
+    with pytest.raises(errors.ClientError):
+      client.models.count_tokens(
+          model=_COUNT_TOKENS_PARAMS.model,
+          contents=_COUNT_TOKENS_PARAMS.contents,
+          config=config,
+      )

google/genai/tests/models/test_delete.py

@@ -0,0 +1,107 @@
+# 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 models.get."""
+
+import pytest
+
+from ... import errors
+from ... import types
+from .. import pytest_helper
+
+TEST_API_VERSION = 'test_api_version'
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_delete_model',
+        parameters=types._DeleteModelParameters(
+            model='models/8533706666867163136'
+        ),
+        exception_if_mldev='404',
+        skip_in_api_mode='cannot be deleted for multiple times',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_delete_model_with_http_options_in_method',
+        parameters=types._DeleteModelParameters(
+            model='tmodels/8533706666867163136',
+            config={
+                'http_options': {
+                    'api_version': TEST_API_VERSION,
+                    'headers': {'test': 'headers'},
+                },
+            },
+        ),
+        exception_if_vertex=TEST_API_VERSION,
+        exception_if_mldev='404',
+        skip_in_api_mode='cannot be deleted for multiple times',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_delete_tuned_model',
+        parameters=types._DeleteModelParameters(
+            model='tunedModels/generate-num-9598'
+        ),
+        exception_if_vertex='404',
+        skip_in_api_mode='cannot be deleted for multiple times',
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.delete',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async_delete_model_with_http_options_in_method(client):
+
+  with pytest.raises(errors.ClientError) as e:
+    await client.aio.models.delete(
+        model='tunedModels/generate-num-888',
+        config={
+            'http_options': {
+                'api_version': TEST_API_VERSION,
+                'headers': {'test': 'headers'},
+            },
+        },
+    )
+  if client._api_client.vertexai:
+    assert TEST_API_VERSION in e.value.args[0]
+  else:
+    assert '404' in str(e)
+
+
+@pytest.mark.asyncio
+async def test_async_delete_tuned_model(client):
+  if client._api_client.vertexai:
+    with pytest.raises(errors.ClientError) as e:
+      await client.aio.models.delete(model='tunedModels/generate-num-888')
+    assert '404' in str(e)
+  else:
+    response = await client.aio.models.delete(
+        model='tunedModels/generate-num-888'
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_delete_model(client):
+  if client._api_client.vertexai:
+    response = await client.aio.models.delete(
+        model='models/1071206899942162432'
+    )
+  else:
+    with pytest.raises(errors.ClientError) as e:
+      await client.aio.models.delete(model='models/1071206899942162432')
+    assert '404' in str(e)

google/genai/tests/models/test_edit_image.py

@@ -0,0 +1,264 @@
+# 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 edit_image."""
+
+import os
+
+import pydantic
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+CAPABILITY_MODEL_NAME = 'imagen-3.0-capability-001'
+
+IMAGE_FILE_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/google.png')
+)
+
+MASK_FILE_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/checkerboard.png')
+)
+
+BRIDGE_IMAGE_FILE_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/bridge1.png')
+)
+
+raw_ref_image = types.RawReferenceImage(
+    reference_id=1,
+    reference_image=types.Image.from_file(location=IMAGE_FILE_PATH),
+)
+
+mask_ref_image = types.MaskReferenceImage(
+    reference_id=2,
+    config=types.MaskReferenceConfig(
+        mask_mode='MASK_MODE_BACKGROUND',
+        mask_dilation=0.06,
+    ),
+)
+
+mask_ref_image_user_provided = types.MaskReferenceImage(
+    reference_id=2,
+    reference_image=types.Image.from_file(location=MASK_FILE_PATH),
+    config=types.MaskReferenceConfig(
+        mask_mode='MASK_MODE_USER_PROVIDED',
+        mask_dilation=0.06,
+    ),
+)
+
+control_ref_image = types.ControlReferenceImage(
+    reference_id=2,
+    reference_image=types.Image.from_file(location=MASK_FILE_PATH),
+    config=types.ControlReferenceConfig(
+        control_type='CONTROL_TYPE_SCRIBBLE',
+        # Backend creates the control image if this is set to True.
+        enable_control_image_computation=False,
+    ),
+)
+
+style_ref_image_customization = types.StyleReferenceImage(
+    reference_id=1,
+    reference_image=types.Image.from_file(location=IMAGE_FILE_PATH),
+    config=types.StyleReferenceConfig(
+        style_description='glowing style',
+    ),
+)
+
+subject_ref_image_customization = types.SubjectReferenceImage(
+    reference_id=1,
+    reference_image=types.Image.from_file(location=IMAGE_FILE_PATH),
+    config=types.SubjectReferenceConfig(
+        subject_type='SUBJECT_TYPE_PRODUCT',
+        subject_description='A product logo that is a multi-colored letter G',
+    ),
+)
+
+dog_content_ref_image = types.ContentReferenceImage(
+    reference_id=1,
+    reference_image=types.Image(
+        gcs_uri='gs://genai-sdk-tests/inputs/images/dog.jpg'
+    ),
+)
+
+cyberpunk_style_ref_image = types.StyleReferenceImage(
+    reference_id=2,
+    reference_image=types.Image(
+        gcs_uri='gs://genai-sdk-tests/inputs/images/cyberpunk.jpg'
+    ),
+    config=types.StyleReferenceConfig(
+        style_description='cyberpunk style',
+    ),
+)
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_edit_mask_inpaint_insert',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._EditImageParameters(
+            model=CAPABILITY_MODEL_NAME,
+            prompt='Sunlight and clear weather',
+            reference_images=[raw_ref_image, mask_ref_image],
+            config=types.EditImageConfig(
+                edit_mode=types.EditMode.EDIT_MODE_INPAINT_INSERTION,
+                number_of_images=1,
+                # Test comprehensive configs
+                # aspect_ratio is not supported for mask editing
+                negative_prompt='human',
+                guidance_scale=15.0,
+                safety_filter_level=types.SafetyFilterLevel.BLOCK_MEDIUM_AND_ABOVE,
+                person_generation=types.PersonGeneration.DONT_ALLOW,
+                include_safety_attributes=False,
+                include_rai_reason=True,
+                output_mime_type='image/jpeg',
+                output_compression_quality=80,
+                base_steps=32,
+                add_watermark=False,
+                labels={'imagen_label_key': 'edit_image'}
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_edit_mask_inpaint_insert_user_provided',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._EditImageParameters(
+            model=CAPABILITY_MODEL_NAME,
+            prompt='Change the colors',
+            reference_images=[raw_ref_image, mask_ref_image_user_provided],
+            config=types.EditImageConfig(
+                edit_mode=types.EditMode.EDIT_MODE_INPAINT_INSERTION,
+                # aspect_ratio is not supported for mask editing
+                number_of_images=1,
+                include_rai_reason=True,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_edit_control_user_provided',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._EditImageParameters(
+            model=CAPABILITY_MODEL_NAME,
+            prompt='Change the colors aligning with the scribble map [2]',
+            reference_images=[raw_ref_image, control_ref_image],
+            config=types.EditImageConfig(
+                number_of_images=1,
+                aspect_ratio='9:16',
+                include_rai_reason=True,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_edit_style_reference_image_customization',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._EditImageParameters(
+            model=CAPABILITY_MODEL_NAME,
+            prompt=(
+                'Generate an image in glowing style [1] based on the following'
+                ' caption: A church in the mountain.'
+            ),
+            reference_images=[style_ref_image_customization],
+            config=types.EditImageConfig(
+                number_of_images=1,
+                aspect_ratio='9:16',
+                include_rai_reason=True,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_edit_subject_image_customization',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._EditImageParameters(
+            model=CAPABILITY_MODEL_NAME,
+            prompt=(
+                'Generate an image containing a mug with the product logo [1]'
+                ' visible on the side of the mug.'
+            ),
+            reference_images=[subject_ref_image_customization],
+            config=types.EditImageConfig(
+                number_of_images=1,
+                aspect_ratio='9:16',
+                include_rai_reason=True,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_edit_content_image_ingredients',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._EditImageParameters(
+            model='imagen-4.0-ingredients-preview',
+            prompt=(
+                'Dog in [1] sleeping on the ground at the bottom of the image'
+                ' with the cyberpunk city landscape in [2] in the background.'
+            ),
+            reference_images=[dog_content_ref_image, cyberpunk_style_ref_image],
+            config=types.EditImageConfig(
+                number_of_images=1,
+                aspect_ratio='9:16',
+                include_rai_reason=True,
+                output_mime_type='image/jpeg',
+            ),
+        ),
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.edit_image',
+    test_table=test_table,
+)
+
+
+def test_setting_reference_type_raises(client):
+  with pytest.raises(pydantic.ValidationError):
+    types.SubjectReferenceImage(
+        reference_id=1,
+        # This test should fail because the user can't set reference_type.
+        reference_type='REFERENCE_TYPE_SUBJECT',
+        reference_image=types.Image.from_file(location=IMAGE_FILE_PATH),
+        config=types.SubjectReferenceConfig(
+            subject_type='SUBJECT_TYPE_PRODUCT',
+            subject_description=(
+                'A product logo that is a multi-colored letter G'
+            ),
+        ),
+    )
+
+
+@pytest.mark.asyncio
+async def test_edit_mask_inpaint_insert_async(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    response = await client.aio.models.edit_image(
+        model=CAPABILITY_MODEL_NAME,
+        prompt='Sunlight and clear weather',
+        reference_images=[raw_ref_image, mask_ref_image],
+        config=types.EditImageConfig(
+            edit_mode=types.EditMode.EDIT_MODE_INPAINT_INSERTION,
+            number_of_images=1,
+            # Test comprehensive configs
+            # aspect_ratio is not supported for mask editing
+            negative_prompt='human',
+            guidance_scale=15.0,
+            safety_filter_level=types.SafetyFilterLevel.BLOCK_MEDIUM_AND_ABOVE,
+            person_generation=types.PersonGeneration.DONT_ALLOW,
+            include_safety_attributes=False,
+            include_rai_reason=True,
+            output_mime_type='image/jpeg',
+            output_compression_quality=80,
+            add_watermark=False,
+        ),
+    )
+    assert response.generated_images[0].image.image_bytes