google-genai 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

google/genai/_transformers.py

@@ -17,20 +17,28 @@
 
 import base64
 from collections.abc import Iterable, Mapping
+from enum import Enum, EnumMeta
 import inspect
 import io
 import re
 import time
 import typing
 from typing import Any, GenericAlias, Optional, Union
+import sys
+
+if typing.TYPE_CHECKING:
+  import PIL.Image
 
-import PIL.Image
-import PIL.PngImagePlugin
 import pydantic
 
 from . import _api_client
 from . import types
 
+if sys.version_info >= (3, 11):
+  from types import UnionType
+else:
+  UnionType = typing._UnionGenericAlias
+
 
 def _resource_name(
     client: _api_client.ApiClient,
@@ -192,9 +200,15 @@ def t_caches_model(api_client: _api_client.ApiClient, model: str):
     return model
 
 
-def pil_to_blob(img):
+def pil_to_blob(img) -> types.Blob:
+  try:
+    import PIL.PngImagePlugin
+    PngImagePlugin = PIL.PngImagePlugin
+  except ImportError:
+    PngImagePlugin = None
+
   bytesio = io.BytesIO()
-  if isinstance(img, PIL.PngImagePlugin.PngImageFile) or img.mode == 'RGBA':
+  if PngImagePlugin is not None and isinstance(img, PngImagePlugin.PngImageFile) or img.mode == 'RGBA':
     img.save(bytesio, format='PNG')
     mime_type = 'image/png'
   else:
@@ -205,20 +219,26 @@ def pil_to_blob(img):
   return types.Blob(mime_type=mime_type, data=data)
 
 
-PartType = Union[types.Part, types.PartDict, str, PIL.Image.Image]
+PartType = Union[types.Part, types.PartDict, str, 'PIL.Image.Image']
 
 
 def t_part(client: _api_client.ApiClient, part: PartType) -> types.Part:
+  try:
+    import PIL.Image
+    PIL_Image = PIL.Image.Image
+  except ImportError:
+    PIL_Image = None
+
   if not part:
     raise ValueError('content part is required.')
   if isinstance(part, str):
     return types.Part(text=part)
-  if isinstance(part, PIL.Image.Image):
+  if PIL_Image is not None and isinstance(part, PIL_Image):
     return types.Part(inline_data=pil_to_blob(part))
   if isinstance(part, types.File):
     if not part.uri or not part.mime_type:
       raise ValueError('file uri and mime_type are required.')
-    return types.Part.from_uri(part.uri, part.mime_type)
+    return types.Part.from_uri(file_uri=part.uri, mime_type=part.mime_type)
   else:
     return part
 
@@ -297,104 +317,192 @@ def t_contents(
     return [t_content(client, contents)]
 
 
-def process_schema(
-    data: dict[str, Any], client: Optional[_api_client.ApiClient] = None
-):
-  if isinstance(data, dict):
-    # Iterate over a copy of keys to allow deletion
-    for key in list(data.keys()):
-      # Only delete 'title'for the Gemini API
-      if client and not client.vertexai and key == 'title':
-        del data[key]
-      else:
-        process_schema(data[key], client)
-  elif isinstance(data, list):
-    for item in data:
-      process_schema(item, client)
+def handle_null_fields(schema: dict[str, Any]):
+  """Process null fields in the schema so it is compatible with OpenAPI.
 
-  return data
+  The OpenAPI spec does not support 'type: 'null' in the schema. This function
+  handles this case by adding 'nullable: True' to the null field and removing
+  the {'type': 'null'} entry.
 
+  https://swagger.io/docs/specification/v3_0/data-models/data-types/#null
 
-def _build_schema(fname: str, fields_dict: dict[str, Any]) -> dict[str, Any]:
-  parameters = pydantic.create_model(fname, **fields_dict).model_json_schema()
-  defs = parameters.pop('$defs', {})
+  Example of schema properties before and after handling null fields:
+    Before:
+      {
+        "name": {
+          "title": "Name",
+          "type": "string"
+        },
+        "total_area_sq_mi": {
+          "anyOf": [
+            {
+              "type": "integer"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Total Area Sq Mi"
+        }
+      }
 
-  for _, value in defs.items():
-    unpack_defs(value, defs)
+    After:
+      {
+        "name": {
+          "title": "Name",
+          "type": "string"
+        },
+        "total_area_sq_mi": {
+          "type": "integer",
+          "nullable": true,
+          "default": null,
+          "title": "Total Area Sq Mi"
+        }
+      }
+  """
+  if (
+      isinstance(schema, dict)
+      and 'type' in schema
+      and schema['type'] == 'null'
+  ):
+    schema['nullable'] = True
+    del schema['type']
+  elif 'anyOf' in schema:
+    for item in schema['anyOf']:
+      if 'type' in item and item['type'] == 'null':
+        schema['nullable'] = True
+        schema['anyOf'].remove({'type': 'null'})
+        if len(schema['anyOf']) == 1:
+          # If there is only one type left after removing null, remove the anyOf field.
+          field_type = schema['anyOf'][0]['type']
+          schema['type'] = field_type
+          del schema['anyOf']
 
-  unpack_defs(parameters, defs)
-  return parameters['properties']['dummy']
 
+def process_schema(
+    schema: dict[str, Any],
+    client: Optional[_api_client.ApiClient] = None,
+    defs: Optional[dict[str, Any]]=None):
+  """Updates the schema and each sub-schema inplace to be API-compatible.
 
-def unpack_defs(schema: dict[str, Any], defs: dict[str, Any]):
-  """Unpacks the $defs values in the schema generated by pydantic so they can be understood by the API.
+  - Removes the `title` field from the schema if the client is not vertexai.
+  - Inlines the $defs.
 
-  Example of a schema before and after unpacking:
+  Example of a schema before and after (with mldev):
     Before:
 
     `schema`
 
-    {'properties': {
-        'dummy': {
-            'items': {
-                '$ref': '#/$defs/CountryInfo'
-            },
-            'title': 'Dummy',
-            'type': 'array'
-            }
+    {
+        'items': {
+            '$ref': '#/$defs/CountryInfo'
         },
-        'required': ['dummy'],
-        'title': 'dummy',
-        'type': 'object'}
+        'title': 'Placeholder',
+        'type': 'array'
+    }
+
 
     `defs`
 
-    {'CountryInfo': {'properties': {'continent': {'title': 'Continent', 'type':
-    'string'}, 'gdp': {'title': 'Gdp', 'type': 'integer'}}, 'required':
-    ['continent', 'gdp'], 'title': 'CountryInfo', 'type': 'object'}}
+    {
+      'CountryInfo': {
+        'properties': {
+          'continent': {
+              'title': 'Continent',
+              'type': 'string'
+          },
+          'gdp': {
+              'title': 'Gdp',
+              'type': 'integer'}
+          },
+        }
+        'required':['continent', 'gdp'],
+        'title': 'CountryInfo',
+        'type': 'object'
+      }
+    }
 
     After:
 
     `schema`
-    {'properties': {
-        'continent': {'title': 'Continent', 'type': 'string'},
-        'gdp': {'title': 'Gdp', 'type': 'integer'}
-      },
-      'required': ['continent', 'gdp'],
-      'title': 'CountryInfo',
-      'type': 'object'
+     {
+        'items': {
+          'properties': {
+            'continent': {
+                'type': 'string'
+            },
+            'gdp': {
+                'type': 'integer'}
+            },
+          }
+          'required':['continent', 'gdp'],
+          'type': 'object'
+        },
+        'type': 'array'
     }
   """
-  properties = schema.get('properties', None)
-  if properties is None:
+  if client and not client.vertexai:
+    schema.pop('title', None)
+
+  if defs is None:
+    defs = schema.pop('$defs', {})
+    for _, sub_schema in defs.items():
+      process_schema(sub_schema, client, defs)
+
+  handle_null_fields(schema)
+
+  any_of = schema.get('anyOf', None)
+  if any_of is not None:
+    for sub_schema in any_of:
+      process_schema(sub_schema, client, defs)
     return
 
-  for name, value in properties.items():
-    ref_key = value.get('$ref', None)
-    if ref_key is not None:
-      ref = defs[ref_key.split('defs/')[-1]]
-      unpack_defs(ref, defs)
-      properties[name] = ref
-      continue
-
-    anyof = value.get('anyOf', None)
-    if anyof is not None:
-      for i, atype in enumerate(anyof):
-        ref_key = atype.get('$ref', None)
-        if ref_key is not None:
-          ref = defs[ref_key.split('defs/')[-1]]
-          unpack_defs(ref, defs)
-          anyof[i] = ref
-      continue
-
-    items = value.get('items', None)
-    if items is not None:
-      ref_key = items.get('$ref', None)
-      if ref_key is not None:
+  schema_type = schema.get('type', None)
+  if isinstance(schema_type, Enum):
+    schema_type = schema_type.value
+  schema_type = schema_type.upper()
+
+  if schema_type == 'OBJECT':
+    properties = schema.get('properties', None)
+    if properties is None:
+      return
+    for name, sub_schema in properties.items():
+      ref_key = sub_schema.get('$ref', None)
+      if ref_key is None:
+        process_schema(sub_schema, client, defs)
+      else:
         ref = defs[ref_key.split('defs/')[-1]]
-        unpack_defs(ref, defs)
-        value['items'] = ref
-        continue
+        process_schema(ref, client, defs)
+        properties[name] = ref
+  elif schema_type == 'ARRAY':
+    sub_schema = schema.get('items', None)
+    if sub_schema is None:
+      return
+    ref_key = sub_schema.get('$ref', None)
+    if ref_key is None:
+      process_schema(sub_schema, client, defs)
+    else:
+      ref = defs[ref_key.split('defs/')[-1]]
+      process_schema(ref, client, defs)
+      schema['items'] = ref
+
+def _process_enum(
+    enum: EnumMeta, client: Optional[_api_client.ApiClient] = None
+) -> types.Schema:
+  for member in enum:
+    if not isinstance(member.value, str):
+      raise TypeError(
+          f'Enum member {member.name} value must be a string, got'
+          f' {type(member.value)}'
+      )
+  class Placeholder(pydantic.BaseModel):
+    placeholder: enum
+
+  enum_schema = Placeholder.model_json_schema()
+  process_schema(enum_schema, client)
+  enum_schema = enum_schema['properties']['placeholder']
+  return types.Schema.model_validate(enum_schema)
 
 
 def t_schema(
@@ -403,28 +511,39 @@ def t_schema(
   if not origin:
     return None
   if isinstance(origin, dict):
-    return process_schema(origin, client)
+    process_schema(origin, client)
+    return types.Schema.model_validate(origin)
+  if isinstance(origin, EnumMeta):
+    return _process_enum(origin, client)
   if isinstance(origin, types.Schema):
     if dict(origin) == dict(types.Schema()):
       # response_schema value was coerced to an empty Schema instance because it did not adhere to the Schema field annotation
       raise ValueError(f'Unsupported schema type.')
-    schema = process_schema(origin.model_dump(exclude_unset=True), client)
+    schema = origin.model_dump(exclude_unset=True)
+    process_schema(schema, client)
     return types.Schema.model_validate(schema)
-  if isinstance(origin, GenericAlias):
-    if origin.__origin__ is list:
-      if isinstance(origin.__args__[0], typing.types.UnionType):
-        raise ValueError(f'Unsupported schema type: GenericAlias {origin}')
-      if issubclass(origin.__args__[0], pydantic.BaseModel):
-        # Handle cases where response schema is `list[pydantic.BaseModel]`
-        list_schema = _build_schema(
-            'dummy', {'dummy': (origin, pydantic.Field())}
-        )
-        list_schema = process_schema(list_schema, client)
-        return types.Schema.model_validate(list_schema)
-    raise ValueError(f'Unsupported schema type: GenericAlias {origin}')
-  if issubclass(origin, pydantic.BaseModel):
-    schema = process_schema(origin.model_json_schema(), client)
+
+  if (
+      # in Python 3.9 Generic alias list[int] counts as a type,
+      # and breaks issubclass because it's not a class.
+      not isinstance(origin, GenericAlias) and 
+      isinstance(origin, type) and 
+      issubclass(origin, pydantic.BaseModel)
+  ):
+    schema = origin.model_json_schema()
+    process_schema(schema, client)
     return types.Schema.model_validate(schema)
+  elif (
+        isinstance(origin, GenericAlias) or isinstance(origin, type) or isinstance(origin, UnionType)
+  ):
+    class Placeholder(pydantic.BaseModel):
+      placeholder: origin
+
+    schema = Placeholder.model_json_schema()
+    process_schema(schema, client)
+    schema = schema['properties']['placeholder']
+    return types.Schema.model_validate(schema)
+
   raise ValueError(f'Unsupported schema type: {origin}')
 
 
@@ -464,7 +583,9 @@ def t_tool(client: _api_client.ApiClient, origin) -> types.Tool:
   if inspect.isfunction(origin) or inspect.ismethod(origin):
     return types.Tool(
         function_declarations=[
-            types.FunctionDeclaration.from_callable(client, origin)
+            types.FunctionDeclaration.from_callable(
+                client=client, callable=origin
+            )
         ]
     )
   else: