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

google/genai/tests/models/test_get.py

@@ -0,0 +1,104 @@
+# 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_http_options = {'api_version': 'v1', 'headers': {'test': 'headers'}}
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_get_vertex_tuned_model',
+        parameters=types._GetModelParameters(
+            model='models/2171259487439028224'
+        ),
+        exception_if_mldev='404',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_mldev_tuned_model',
+        parameters=types._GetModelParameters(
+            model='tunedModels/generatenum5443-ekrw7ie9wis23zbeogbw6jq8'
+        ),
+        exception_if_vertex='404',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_vertex_tuned_model_with_http_options_in_method',
+        parameters=types._GetModelParameters(
+            model='models/2171259487439028224',
+            config={
+                'http_options': test_http_options,
+            },
+        ),
+        exception_if_mldev='404',
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_mldev_base_model_with_http_options_in_method',
+        parameters=types._GetModelParameters(
+            model='gemini-2.5-flash',
+            config={
+                'http_options': test_http_options,
+            },
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_base_model',
+        parameters=types._GetModelParameters(model='gemini-2.5-flash'),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_get_base_model_with_models_prefix',
+        parameters=types._GetModelParameters(model='models/gemini-2.5-flash'),
+        exception_if_vertex='400',
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.get',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_async_get_tuned_model(client):
+  if client._api_client.vertexai:
+    with pytest.raises(errors.ClientError) as e:
+      await client.aio.models.get(model='tunedModels/generate-num-1896')
+    assert '404' in str(e)
+  else:
+    response = await client.aio.models.get(
+        model='tunedModels/generate-num-1896'
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_get_model(client):
+  if client._api_client.vertexai:
+    response = await client.aio.models.get(
+        model='models/7687416965014487040',
+        config={'http_options': test_http_options},
+    )
+  else:
+    with pytest.raises(errors.ClientError) as e:
+      await client.aio.models.get(
+          model='models/7687416965014487040',
+          config={'http_options': test_http_options},
+      )
+    assert '404' in str(e)

google/genai/tests/models/test_list.py

@@ -0,0 +1,233 @@
+# 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.list."""
+
+from unittest import mock
+
+import pytest
+
+from ... import client as genai_client
+from ... import _api_client
+from ... import types
+from .. import pytest_helper
+
+test_http_options = {'headers': {'test': 'headers'}}
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_tuned_models',
+        parameters=types._ListModelsParameters(
+            config={'query_base': False, 'page_size': 3}
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_tuned_models_with_filter',
+        parameters=types._ListModelsParameters(
+            config={'query_base': False,
+                    'page_size': 3,
+                    'filter': 'displayName="gemini-2.5-flash-1b0689e3-9773-43b4-97eb-c8140d5f183b"'}
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_base_models',
+        parameters=types._ListModelsParameters(),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_base_models_with_config',
+        parameters=types._ListModelsParameters(config={'query_base': True, 'page_size': 10}),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_with_config',
+        parameters=types._ListModelsParameters(config={'page_size': 3}),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_list_models_with_http_options_in_method',
+        parameters=types._ListModelsParameters(
+            config={'page_size': 3, 'http_options': test_http_options}
+        ),
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.list',
+    test_table=test_table,
+)
+
+
+@pytest.fixture()
+def mock_api_client():
+  api_client = mock.MagicMock(spec=genai_client.BaseApiClient)
+  api_client.api_key = 'fake_api_key'
+  api_client._host = lambda: 'fake_host'
+  api_client._http_options = {'headers': {}}
+  api_client.vertexai = False
+  return api_client
+
+
+def test_tuned_models_pager(client):
+  pager = client.models.list(config={'page_size': 10})
+
+  assert pager.name == 'models'
+  assert pager.page_size == 10
+  assert len(pager) <= 10
+
+  # Iterate through all the pages. Then next_page() should raise an exception.
+  for _ in pager:
+    pass
+  with pytest.raises(IndexError, match='No more pages to fetch.'):
+    pager.next_page()
+
+
+def test_base_models_pager(client):
+  pager = client.models.list(config={'page_size': 10, 'query_base': True})
+
+  assert pager.name == 'models'
+  assert pager.page_size == 10
+  assert len(pager) <= 10
+
+  # Iterate through all the pages. Then next_page() should raise an exception.
+  for _ in pager:
+    assert 'content-type' in pager.sdk_http_response.headers
+    assert 'content-encoding' in pager.sdk_http_response.headers
+    pass
+  with pytest.raises(IndexError, match='No more pages to fetch.'):
+    pager.next_page()
+
+
+def test_base_response_with_empty_json_payload_and_http_headers(
+    mock_api_client, client
+):
+  with mock.patch.object(
+      genai_client.Client, '_get_api_client'
+  ) as patch_api_client:
+    patch_api_client.return_value = mock_api_client
+    mock_client = genai_client.Client()
+    sdk_http_response = types.HttpResponse(
+        headers={'header_key': 'header_value'},
+        body='{}',
+    )
+    with mock.patch.object(
+        mock_api_client, 'request', return_value=sdk_http_response
+    ):
+      pager = mock_client.models.list()
+
+      assert len(pager) == 0
+
+
+def test_unknown_json_payload(mock_api_client, client):
+  with mock.patch.object( 
+      genai_client.Client, '_get_api_client'
+  ) as patch_api_client:
+    patch_api_client.return_value = mock_api_client
+    mock_client = genai_client.Client()
+    sdk_http_response = types.HttpResponse(
+        headers={'header_key': 'header_value'},
+        body='{"unknown_key": "unknown_value"}',
+    )
+    with mock.patch.object(
+        mock_api_client, 'request', return_value=sdk_http_response
+    ) :
+      pager = mock_client.models.list()
+
+      assert len(pager) == 0
+
+
+def test_empty_json_payload(mock_api_client, client):
+  with mock.patch.object( 
+      genai_client.Client, '_get_api_client'
+  ) as patch_api_client:
+    patch_api_client.return_value = mock_api_client
+    mock_client = genai_client.Client()
+    sdk_http_response = types.HttpResponse(
+        headers={'header_key': 'header_value'},
+        body='',
+    )
+    with mock.patch.object(
+        mock_api_client, 'request', return_value=sdk_http_response
+    ) :
+      pager = mock_client.models.list()
+
+      assert len(pager) == 0
+
+
+def test_empty_api_response_none_headers(mock_api_client, client):
+  with mock.patch.object(
+      genai_client.Client, '_get_api_client'
+  ) as patch_api_client:
+    patch_api_client.return_value = mock_api_client
+    mock_client = genai_client.Client()
+    sdk_http_response = types.HttpResponse(
+        headers=None,
+        body='{}',
+    )
+    with mock.patch.object(
+        mock_api_client, 'request', return_value=sdk_http_response
+    ):
+      pager = mock_client.models.list()
+
+      assert len(pager) == 0
+
+
+def test_empty_api_response_empty_dict_headers(mock_api_client, client):
+  with mock.patch.object(
+      genai_client.Client, '_get_api_client'
+  ) as patch_api_client:
+    patch_api_client.return_value = mock_api_client
+    mock_client = genai_client.Client()
+    sdk_http_response = types.HttpResponse(
+        headers={},
+        body='{}',
+    )
+    with mock.patch.object(
+        mock_api_client, 'request', return_value=sdk_http_response
+    ):
+      pager = mock_client.models.list()
+
+      assert len(pager) == 0
+
+
+@pytest.mark.asyncio
+async def test_tuned_models_async_pager(client):
+  pager = await client.aio.models.list(config={'page_size': 3, 'query_base': False})
+
+  assert 'Content-Type' in pager.sdk_http_response.headers
+  assert 'Content-Encoding' in pager.sdk_http_response.headers
+  assert pager.name == 'models'
+  assert pager.page_size == 3
+  assert len(pager) <= 3
+
+  # Iterate through all the pages. Then next_page() should raise an exception.
+  async for _ in pager:
+    pass
+  with pytest.raises(IndexError, match='No more pages to fetch.'):
+    await pager.next_page()
+
+
+@pytest.mark.asyncio
+async def test_base_models_async_pager(client):
+  pager = await client.aio.models.list(config={'page_size': 10})
+
+  assert pager.name == 'models'
+  assert pager.page_size == 10
+  assert len(pager) <= 10
+
+  # Iterate through all the pages. Then next_page() should raise an exception.
+  async for _ in pager:
+    pass
+  with pytest.raises(IndexError, match='No more pages to fetch.'):
+    await pager.next_page()

google/genai/tests/models/test_recontext_image.py

@@ -0,0 +1,189 @@
+# 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 recontext_image."""
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+PRODUCT_RECONTEXT_MODEL_LATEST = 'imagen-product-recontext-preview-06-30'
+
+VIRTUAL_TRY_ON_IMAGE_MODEL_LATEST = 'virtual-try-on-preview-08-04'
+
+BACKPACK1 = types.Image(
+    gcs_uri='gs://genai-sdk-tests/inputs/images/backpack1.png',
+)
+
+BACKPACK2 = types.Image(
+    gcs_uri='gs://genai-sdk-tests/inputs/images/backpack2.png',
+)
+
+BACKPACK3 = types.Image(
+    gcs_uri='gs://genai-sdk-tests/inputs/images/backpack3.png',
+)
+
+PERSON_IMAGE = types.Image(
+    gcs_uri='gs://genai-sdk-tests/inputs/images/man.jpg',
+)
+
+PANTS = types.Image(
+    gcs_uri='gs://genai-sdk-tests/inputs/images/pants.jpg',
+)
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_product_recontext',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._RecontextImageParameters(
+            model=PRODUCT_RECONTEXT_MODEL_LATEST,
+            source=types.RecontextImageSource(
+                prompt='On a school desk',
+                product_images=[
+                    types.ProductImage(product_image=BACKPACK1)
+                ],
+            ),
+            config=types.RecontextImageConfig(
+                number_of_images=1,
+                add_watermark=True,
+                output_mime_type='image/jpeg',
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_product_recontext_all_config',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._RecontextImageParameters(
+            model=PRODUCT_RECONTEXT_MODEL_LATEST,
+            source=types.RecontextImageSource(
+                prompt='On a school desk',
+                product_images=[
+                    types.ProductImage(
+                        product_image=BACKPACK1,
+                    ),
+                    types.ProductImage(
+                        product_image=BACKPACK2,
+                    ),
+                    types.ProductImage(
+                        product_image=BACKPACK3,
+                    ),
+                ],
+            ),
+            config=types.RecontextImageConfig(
+                number_of_images=1,
+                output_mime_type='image/jpeg',
+                base_steps=32,
+                seed=1337,
+                add_watermark=False,
+                safety_filter_level=types.SafetyFilterLevel.BLOCK_MEDIUM_AND_ABOVE,
+                person_generation=types.PersonGeneration.ALLOW_ADULT,
+                output_compression_quality=75,
+                enhance_prompt=False,
+                labels={'imagen_label_key': 'recontext_image'}
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_virtual_try_on',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._RecontextImageParameters(
+            model=VIRTUAL_TRY_ON_IMAGE_MODEL_LATEST,
+            source=types.RecontextImageSource(
+                person_image=PERSON_IMAGE,
+                product_images=[
+                    types.ProductImage(product_image=PANTS)
+                ],
+            ),
+            config=types.RecontextImageConfig(
+                number_of_images=1,
+                output_mime_type='image/jpeg',
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_virtual_try_on_all_config',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._RecontextImageParameters(
+            model=VIRTUAL_TRY_ON_IMAGE_MODEL_LATEST,
+            source=types.RecontextImageSource(
+                person_image=PERSON_IMAGE,
+                product_images=[
+                    types.ProductImage(
+                        product_image=PANTS
+                    )
+                ],
+            ),
+            config=types.RecontextImageConfig(
+                number_of_images=1,
+                output_mime_type='image/jpeg',
+                base_steps=32,
+                seed=1337,
+                add_watermark=False,
+                safety_filter_level=types.SafetyFilterLevel.BLOCK_MEDIUM_AND_ABOVE,
+                person_generation=types.PersonGeneration.ALLOW_ADULT,
+                output_compression_quality=75,
+                enhance_prompt=False,
+            ),
+        ),
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.recontext_image',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_product_recontext_async(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    response = await client.aio.models.recontext_image(
+        model=PRODUCT_RECONTEXT_MODEL_LATEST,
+        source=types.RecontextImageSource(
+            prompt='On a school desk',
+            product_images=[
+                types.ProductImage(product_image=BACKPACK1)
+            ],
+        ),
+        config=types.RecontextImageConfig(
+            number_of_images=1,
+            output_mime_type='image/jpeg',
+        )
+    )
+    assert len(response.generated_images) == 1
+    assert response.generated_images[0].image.image_bytes
+
+
+@pytest.mark.asyncio
+async def test_virtual_try_on_async(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    response = await client.aio.models.recontext_image(
+        model=VIRTUAL_TRY_ON_IMAGE_MODEL_LATEST,
+        source=types.RecontextImageSource(
+            person_image=PERSON_IMAGE,
+            product_images=[
+                types.ProductImage(product_image=PANTS)
+            ],
+        ),
+        config=types.RecontextImageConfig(
+            number_of_images=1,
+            output_mime_type='image/jpeg',
+        ),
+    )
+    assert len(response.generated_images) == 1
+    assert response.generated_images[0].image.image_bytes

google/genai/tests/models/test_segment_image.py

@@ -0,0 +1,148 @@
+# 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 segment_image."""
+
+import os
+
+import pytest
+
+from ... import types
+from .. import pytest_helper
+
+SEGMENT_IMAGE_MODEL_LATEST = 'image-segmentation-001'
+
+SOURCE_IMAGE_FILE_PATH1 = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/google.png')
+)
+
+SOURCE_IMAGE1 = types.Image.from_file(location=SOURCE_IMAGE_FILE_PATH1)
+
+SOURCE_IMAGE_FILE_PATH2 = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/skateboard_stop_sign.jpg')
+)
+
+SOURCE_IMAGE2 = types.Image.from_file(location=SOURCE_IMAGE_FILE_PATH2)
+
+SCRIBBLE_IMAGE_FILE_PATH = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), '../data/segmentation_scribble.jpg')
+)
+
+SCRIBBLE_IMAGE = types.Image.from_file(location=SCRIBBLE_IMAGE_FILE_PATH)
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_segment_foreground',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._SegmentImageParameters(
+            model=SEGMENT_IMAGE_MODEL_LATEST,
+            source=types.SegmentImageSource(
+                image=SOURCE_IMAGE1,
+            ),
+            config=types.SegmentImageConfig(
+                mode=types.SegmentMode.FOREGROUND,
+                max_predictions=1,
+                confidence_threshold=0.02,
+                mask_dilation=0.02,
+                binary_color_threshold=98,
+                labels={'imagen_label_key': 'segment_image'}
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_segment_background',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._SegmentImageParameters(
+            model=SEGMENT_IMAGE_MODEL_LATEST,
+            source=types.SegmentImageSource(
+                image=SOURCE_IMAGE1,
+            ),
+            config=types.SegmentImageConfig(
+                mode=types.SegmentMode.BACKGROUND,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_segment_prompt',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._SegmentImageParameters(
+            model=SEGMENT_IMAGE_MODEL_LATEST,
+            source=types.SegmentImageSource(
+                image=SOURCE_IMAGE1,
+                prompt='The letter G',
+            ),
+            config=types.SegmentImageConfig(
+                mode=types.SegmentMode.PROMPT,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_segment_semantic',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._SegmentImageParameters(
+            model=SEGMENT_IMAGE_MODEL_LATEST,
+            source=types.SegmentImageSource(
+                image=SOURCE_IMAGE2,
+                prompt='skateboard',
+            ),
+            config=types.SegmentImageConfig(
+                mode=types.SegmentMode.SEMANTIC,
+            ),
+        ),
+    ),
+    pytest_helper.TestTableItem(
+        name='test_segment_interactive',
+        exception_if_mldev='only supported in the Vertex AI client',
+        parameters=types._SegmentImageParameters(
+            model=SEGMENT_IMAGE_MODEL_LATEST,
+            source=types.SegmentImageSource(
+                image=SOURCE_IMAGE2,
+                scribble_image=types.ScribbleImage(image=SCRIBBLE_IMAGE),
+            ),
+            config=types.SegmentImageConfig(
+                mode=types.SegmentMode.INTERACTIVE,
+            ),
+        ),
+    ),
+]
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='models.segment_image',
+    test_table=test_table,
+)
+
+
+@pytest.mark.asyncio
+async def test_segment_foreground_async(client):
+  with pytest_helper.exception_if_mldev(client, ValueError):
+    response = await client.aio.models.segment_image(
+        model=SEGMENT_IMAGE_MODEL_LATEST,
+        source=types.SegmentImageSource(
+            image=SOURCE_IMAGE2,
+        ),
+        config=types.SegmentImageConfig(
+            mode=types.SegmentMode.FOREGROUND,
+            max_predictions=1,
+            confidence_threshold=0.02,
+            mask_dilation=0.02,
+            binary_color_threshold=98,
+        )
+    )
+    assert response.generated_masks[0].mask.image_bytes
+    assert len(response.generated_masks) == 1
+    assert response.generated_masks[0].labels[0].label == 'foreground'
+    assert response.generated_masks[0].labels[0].score > 0