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

google/genai/tests/transformers/test_schema.py

@@ -0,0 +1,653 @@
+# 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 schema processing methods in the _transformers module."""
+
+import copy
+from typing import Union
+
+import pydantic
+import pytest
+
+from ... import _transformers
+from ... import client as google_genai_client_module
+from ... import types
+
+
+class CurrencyInfo(pydantic.BaseModel):
+  name: str
+  code: str
+  symbol: str
+
+
+currency_info_fields = CurrencyInfo.model_fields
+
+
+class CountryInfo(pydantic.BaseModel):
+  name: str
+  population: int
+  capital: str
+  continent: str
+  gdp: int
+  official_language: str
+  total_area_sq_mi: int
+
+
+country_info_fields = CountryInfo.model_fields
+
+
+class CountryInfoWithCurrency(pydantic.BaseModel):
+  name: str
+  population: int
+  capital: str
+  continent: str
+  gdp: int
+  official_language: str
+  total_area_sq_mi: int
+  currency: CurrencyInfo
+
+
+nested_country_info_fields = CountryInfoWithCurrency.model_fields
+
+
+class CountryInfoWithNullFields(pydantic.BaseModel):
+  name: str
+  population: Union[int, None] = None
+
+
+class CountryInfoWithDefaultValue(pydantic.BaseModel):
+  name: str
+  population: int = 0
+
+
+class CountryInfoWithAnyOf(pydantic.BaseModel):
+  name: str
+  restaurants_per_capita: Union[int, float]
+
+
+@pytest.fixture
+@pytest.mark.parametrize('use_vertex', [True, False])
+def client(use_vertex):
+  if use_vertex:
+    yield google_genai_client_module.Client(
+        vertexai=use_vertex, project='test-project', location='test-location'
+    )
+  else:
+    yield google_genai_client_module.Client(
+        vertexai=use_vertex, api_key='test-api-key'
+    )
+
+
+def test_build_schema_for_list_of_pydantic_schema(client):
+  """Tests _build_schema() when list[pydantic.BaseModel] is provided to response_schema."""
+
+  list_schema = _transformers.t_schema(client, CountryInfo).model_dump()
+
+  assert isinstance(list_schema, dict)
+
+  for field_name in list_schema['properties']:
+    assert 'title' in list_schema['properties'][field_name]
+    assert 'type' in list_schema['properties'][field_name]
+    assert field_name in country_info_fields
+    field_type_str = country_info_fields[field_name].annotation.__name__
+    assert (
+        list_schema['properties'][field_name]['type']
+        .lower()
+        .startswith(field_type_str.lower())
+    )
+    assert 'required' in list_schema
+    assert list_schema['required'] == list(country_info_fields.keys())
+
+
+def test_build_schema_for_list_of_nested_pydantic_schema(client):
+  """Tests _build_schema() when list[pydantic.BaseModel] is provided to response_schema and the pydantic.BaseModel has nested pydantic fields."""
+  list_schema = _transformers.t_schema(
+      client, CountryInfoWithCurrency
+  ).model_dump()
+
+  assert isinstance(list_schema, dict)
+
+  for field_name in list_schema['properties']:
+    assert 'title' in list_schema['properties'][field_name]
+    assert 'type' in list_schema['properties'][field_name]
+    assert field_name in nested_country_info_fields
+
+  # Tested nested schema was created
+  assert 'properties' in list_schema['properties']['currency']
+
+  for field_name in list_schema['properties']['currency']['properties']:
+    assert field_name in currency_info_fields
+
+
+def test_t_schema_for_pydantic_schema(client):
+  """Tests t_schema when pydantic.BaseModel is passed to response_schema."""
+  transformed_schema = _transformers.t_schema(client, CountryInfo)
+  assert isinstance(transformed_schema, types.Schema)
+  for schema_property in transformed_schema.properties:
+    assert schema_property in country_info_fields
+    assert isinstance(
+        transformed_schema.properties[schema_property], types.Schema
+    )
+
+
+def test_t_schema_for_list_of_pydantic_schema(client):
+  """Tests t_schema when list[pydantic.BaseModel] is passed to response_schema."""
+  transformed_schema = _transformers.t_schema(client, list[CountryInfo])
+  assert isinstance(transformed_schema, types.Schema)
+  assert isinstance(transformed_schema.items, types.Schema)
+
+  for schema_property in transformed_schema.items.properties:
+    assert schema_property in country_info_fields
+    assert isinstance(
+        transformed_schema.items.properties[schema_property], types.Schema
+    )
+
+
+def test_t_schema_for_null_fields(client):
+  """Tests t_schema when null fields are present."""
+  transformed_schema = _transformers.t_schema(client, CountryInfoWithNullFields)
+  assert isinstance(transformed_schema, types.Schema)
+  assert transformed_schema.properties['population'].nullable
+
+
+def test_schema_with_no_null_fields_is_unchanged():
+  """Tests handle_null_fields() doesn't change anything when no null fields are present."""
+  test_properties = {
+      'name': {'title': 'Name', 'type': 'string'},
+      'total_area_sq_mi': {
+          'anyOf': [{'type': 'integer'}, {'type': 'float'}],
+          'default': 'null',
+          'title': 'Total Area Sq Mi',
+      },
+  }
+
+  for _, schema in test_properties.items():
+    schema_before = copy.deepcopy(schema)
+    _transformers.handle_null_fields(schema)
+    assert schema_before == schema
+
+
+@pytest.mark.parametrize('use_vertex', [True, False])
+def test_schema_with_default_value(client):
+
+  transformed_schema = _transformers.t_schema(
+      client._api_client, CountryInfoWithDefaultValue
+  )
+  expected_schema = types.Schema(
+      properties={
+          'name': types.Schema(
+              type='STRING',
+              title='Name',
+          ),
+          'population': types.Schema(
+              type='INTEGER',
+              default=0,
+              title='Population',
+          ),
+      },
+      type='OBJECT',
+      required=['name'],
+      title='CountryInfoWithDefaultValue',
+      property_ordering=['name', 'population'],
+  )
+
+  assert transformed_schema == expected_schema
+
+
+def test_schema_with_any_of(client):
+  transformed_schema = _transformers.t_schema(client, CountryInfoWithAnyOf)
+  expected_schema = types.Schema(
+      properties={
+          'name': types.Schema(
+              type='STRING',
+              title='Name',
+          ),
+          'restaurants_per_capita': types.Schema(
+              any_of=[
+                  types.Schema(type='INTEGER'),
+                  types.Schema(type='NUMBER'),
+              ],
+              title='Restaurants Per Capita',
+          ),
+      },
+      type='OBJECT',
+      required=['name', 'restaurants_per_capita'],
+      title='CountryInfoWithAnyOf',
+      property_ordering=['name', 'restaurants_per_capita'],
+  )
+
+  assert transformed_schema == expected_schema
+
+
+@pytest.mark.parametrize('use_vertex', [True, False])
+def test_complex_dict_schema_with_anyof_is_unchanged(client):
+  """When a dict schema is passed to process_schema, the only change should be camel-casing anyOf."""
+  if client.vertexai:
+    dict_schema = {
+        'type': 'OBJECT',
+        'title': 'Fruit Basket',
+        'description': 'A structured representation of a fruit basket',
+        'required': ['fruit'],
+        'properties': {
+            'fruit': {
+                'type': 'ARRAY',
+                'description': 'An ordered list of the fruit in the basket',
+                'items': {
+                    'description': 'A piece of fruit',
+                    'anyOf': [
+                        {
+                            'title': 'Apple',
+                            'description': 'Describes an apple',
+                            'type': 'OBJECT',
+                            'properties': {
+                                'type': {
+                                    'type': 'STRING',
+                                    'description': "Always 'apple'",
+                                },
+                                'color': {
+                                    'type': 'STRING',
+                                    'description': (
+                                        "The color of the apple (e.g., 'red')"
+                                    ),
+                                },
+                            },
+                            'propertyOrdering': ['type', 'color'],
+                            'required': ['type', 'color'],
+                        },
+                        {
+                            'title': 'Orange',
+                            'description': 'Describes an orange',
+                            'type': 'OBJECT',
+                            'properties': {
+                                'type': {
+                                    'type': 'STRING',
+                                    'description': "Always 'orange'",
+                                },
+                                'size': {
+                                    'type': 'STRING',
+                                    'description': (
+                                        'The size of the orange (e.g.,'
+                                        " 'medium')"
+                                    ),
+                                },
+                            },
+                            'propertyOrdering': ['type', 'size'],
+                            'required': ['type', 'size'],
+                        },
+                    ],
+                },
+            }
+        },
+    }
+
+    schema_before = copy.deepcopy(dict_schema)
+    _transformers.process_schema(dict_schema, client)
+
+    assert schema_before == dict_schema
+
+
+@pytest.mark.parametrize('use_vertex', [True, False])
+def test_process_schema_converts_const_to_enum(client):
+  """The 'const' field should be converted to a singleton 'enum'."""
+  schema = {
+      'type': 'STRING',
+      'const': 'FOO',
+  }
+  expected_schema = {
+      'type': 'STRING',
+      'enum': ['FOO'],
+  }
+
+  _transformers.process_schema(schema, client)
+
+  assert schema == expected_schema
+
+
+@pytest.mark.parametrize('use_vertex', [True, False])
+def test_process_schema_forbids_non_string_const(client):
+  """The 'const' field only works for strings."""
+  schema = {
+      'type': 'INTEGER',
+      'const': 123,
+  }
+
+  with pytest.raises(ValueError, match='.*Literal values must be strings.*'):
+    _transformers.process_schema(schema, client)
+
+
+@pytest.mark.parametrize(
+    'use_vertex,order_properties',
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_process_schema_order_properties_propagates_into_defs(
+    client, order_properties
+):
+  """The `order_properties` setting should apply to '$defs'."""
+  schema = {
+      '$ref': '#/$defs/Foo',
+      '$defs': {
+          'Foo': {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+          },
+      },
+  }
+  schema_without_property_ordering = {
+      'type': 'OBJECT',
+      'properties': {
+          'foo': {'type': 'STRING'},
+          'bar': {'type': 'STRING'},
+      },
+  }
+  schema_with_property_ordering = {
+      'type': 'OBJECT',
+      'properties': {
+          'foo': {'type': 'STRING'},
+          'bar': {'type': 'STRING'},
+      },
+      'property_ordering': ['foo', 'bar'],
+  }
+
+  _transformers.process_schema(
+      schema, client, order_properties=order_properties
+  )
+
+  if order_properties:
+    assert schema == schema_with_property_ordering
+  else:
+    assert schema == schema_without_property_ordering
+
+
+@pytest.mark.parametrize(
+    'use_vertex,order_properties',
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_process_schema_order_properties_propagates_into_items(
+    client, order_properties
+):
+  """The `order_properties` setting should apply to 'items'."""
+  schema = {
+      'type': 'ARRAY',
+      'items': {
+          'type': 'OBJECT',
+          'properties': {
+              'foo': {'type': 'STRING'},
+              'bar': {'type': 'STRING'},
+          },
+      },
+  }
+  schema_without_property_ordering = copy.deepcopy(schema)
+  schema_with_property_ordering = {
+      'type': 'ARRAY',
+      'items': {
+          'type': 'OBJECT',
+          'properties': {
+              'foo': {'type': 'STRING'},
+              'bar': {'type': 'STRING'},
+          },
+          'property_ordering': ['foo', 'bar'],
+      },
+  }
+
+  _transformers.process_schema(
+      schema, client, order_properties=order_properties
+  )
+
+  if order_properties:
+    assert schema == schema_with_property_ordering
+  else:
+    assert schema == schema_without_property_ordering
+
+
+@pytest.mark.parametrize(
+    'use_vertex,order_properties',
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_process_schema_order_properties_propagates_into_prefix_items(
+    client, order_properties
+):
+  """The `order_properties` setting should apply to 'prefixItems'."""
+  schema = {
+      'type': 'ARRAY',
+      'prefixItems': [
+          {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+          },
+      ],
+  }
+  schema_without_property_ordering = copy.deepcopy(schema)
+  schema_with_property_ordering = {
+      'type': 'ARRAY',
+      'prefixItems': [
+          {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+              'property_ordering': ['foo', 'bar'],
+          },
+      ],
+  }
+
+  _transformers.process_schema(
+      schema, client, order_properties=order_properties
+  )
+
+  if order_properties:
+    assert schema == schema_with_property_ordering
+  else:
+    assert schema == schema_without_property_ordering
+
+
+@pytest.mark.parametrize(
+    'use_vertex,order_properties',
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_process_schema_order_properties_propagates_into_properties(
+    client, order_properties
+):
+  """The `order_properties` setting should apply to 'properties'."""
+  schema = {
+      'type': 'OBJECT',
+      'properties': {
+          'xyz': {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+          },
+          'abc': {'type': 'STRING'},
+      },
+  }
+  schema_without_property_ordering = copy.deepcopy(schema)
+  schema_with_property_ordering = {
+      'type': 'OBJECT',
+      'properties': {
+          'xyz': {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+              'property_ordering': ['foo', 'bar'],
+          },
+          'abc': {'type': 'STRING'},
+      },
+      'property_ordering': ['xyz', 'abc'],
+  }
+
+  _transformers.process_schema(
+      schema, client, order_properties=order_properties
+  )
+
+  if order_properties:
+    assert schema == schema_with_property_ordering
+  else:
+    assert schema == schema_without_property_ordering
+
+
+@pytest.mark.parametrize(
+    'use_vertex,order_properties',
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_process_schema_order_properties_propagates_into_additional_properties(
+    client, order_properties
+):
+  """The `order_properties` setting should apply to 'additionalProperties'."""
+  schema = {
+      'type': 'OBJECT',
+      'additionalProperties': {
+          'type': 'OBJECT',
+          'properties': {
+              'foo': {'type': 'STRING'},
+              'bar': {'type': 'STRING'},
+          },
+      },
+  }
+  schema_without_property_ordering = copy.deepcopy(schema)
+  schema_with_property_ordering = {
+      'type': 'OBJECT',
+      'additionalProperties': {
+          'type': 'OBJECT',
+          'properties': {
+              'foo': {'type': 'STRING'},
+              'bar': {'type': 'STRING'},
+          },
+          'property_ordering': ['foo', 'bar'],
+      },
+  }
+
+  if client.vertexai:
+    _transformers.process_schema(
+        schema, client, order_properties=order_properties
+    )
+
+    if order_properties:
+      assert schema == schema_with_property_ordering
+    else:
+      assert schema == schema_without_property_ordering
+  else:
+    with pytest.raises(ValueError) as e:
+      _transformers.process_schema(
+          schema, client, order_properties=order_properties
+      )
+    assert 'additionalProperties is not supported in the Gemini API.' in str(e)
+
+
+@pytest.mark.parametrize(
+    'use_vertex,order_properties',
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+def test_process_schema_order_properties_propagates_into_any_of(
+    client, order_properties
+):
+  """The `order_properties` setting should apply to 'anyOf'."""
+  schema = {
+      'anyOf': [
+          {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+          },
+          {'type': 'STRING'},
+      ]
+  }
+  schema_without_property_ordering = copy.deepcopy(schema)
+  schema_with_property_ordering = {
+      'anyOf': [
+          {
+              'type': 'OBJECT',
+              'properties': {
+                  'foo': {'type': 'STRING'},
+                  'bar': {'type': 'STRING'},
+              },
+              'property_ordering': ['foo', 'bar'],
+          },
+          {'type': 'STRING'},
+      ]
+  }
+
+  _transformers.process_schema(
+      schema, client, order_properties=order_properties
+  )
+
+  if order_properties:
+    assert schema == schema_with_property_ordering
+  else:
+    assert schema == schema_without_property_ordering
+
+
+def test_t_schema_does_not_change_property_ordering_if_set(client):
+  """Tests t_schema doesn't overwrite the property_ordering field if already set."""
+
+  schema = CountryInfo.model_json_schema()
+  custom_property_ordering = ['code', 'symbol', 'name']
+  schema['property_ordering'] = custom_property_ordering.copy()
+
+  transformed_schema = _transformers.t_schema(client, schema)
+  assert transformed_schema.property_ordering == custom_property_ordering
+
+
+def test_t_schema_sets_property_ordering_for_json_schema(client):
+  """Tests t_schema sets the property_ordering field for json schemas."""
+
+  schema = CountryInfo.model_json_schema()
+
+  transformed_schema = _transformers.t_schema(client, schema)
+  assert transformed_schema.property_ordering == [
+      'name',
+      'population',
+      'capital',
+      'continent',
+      'gdp',
+      'official_language',
+      'total_area_sq_mi',
+  ]
+
+
+def test_t_schema_sets_property_ordering_for_schema_type(client):
+  """Tests t_schema sets the property_ordering field for Schema types."""
+
+  schema = types.Schema(
+      properties={
+          'name': types.Schema(
+              type='STRING',
+              title='Name',
+          ),
+          'population': types.Schema(
+              type='INTEGER',
+              default=0,
+              title='Population',
+          ),
+      },
+      type='OBJECT',
+      required=['name'],
+      title='CountryInfoWithDefaultValue',
+  )
+
+  transformed_schema = _transformers.t_schema(client, schema)
+  assert transformed_schema.property_ordering == ['name', 'population']