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

google/genai/_transformers.py

@@ -17,25 +17,35 @@
 
 import base64
 from collections.abc import Iterable, Mapping
+from enum import Enum, EnumMeta
 import inspect
 import io
 import re
 import time
-from typing import Any, Optional, Union
+import typing
+from typing import Any, GenericAlias, Optional, Union
+import sys
 
-import PIL.Image
-import PIL.PngImagePlugin
+if typing.TYPE_CHECKING:
+  import PIL.Image
+
+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,
     resource_name: str,
     *,
     collection_identifier: str,
-    collection_hirearchy_depth: int = 2,
+    collection_hierarchy_depth: int = 2,
 ):
   # pylint: disable=line-too-long
   """Prepends resource name with project, location, collection_identifier if needed.
@@ -48,13 +58,13 @@ def _resource_name(
   Args:
     client: The API client.
     resource_name: The user input resource name to be completed.
-    collection_identifier: The collection identifier to be prepended.
-        See collection identifiers in https://google.aip.dev/122.
-    collection_hirearchy_depth: The collection hierarchy depth.
-        Only set this field when the resource has nested collections.
-        For example, `users/vhugo1802/events/birthday-dinner-226`, the
-        collection_identifier is `users` and collection_hirearchy_depth is 4.
-        See nested collections in https://google.aip.dev/122.
+    collection_identifier: The collection identifier to be prepended. See
+      collection identifiers in https://google.aip.dev/122.
+    collection_hierarchy_depth: The collection hierarchy depth. Only set this
+      field when the resource has nested collections. For example,
+      `users/vhugo1802/events/birthday-dinner-226`, the collection_identifier is
+      `users` and collection_hierarchy_depth is 4. See nested collections in
+      https://google.aip.dev/122.
 
   Example:
 
@@ -62,7 +72,8 @@ def _resource_name(
     client.vertexai = True
     client.project = 'bar'
     client.location = 'us-west1'
-    _resource_name(client, 'cachedContents/123', collection_identifier='cachedContents')
+    _resource_name(client, 'cachedContents/123',
+      collection_identifier='cachedContents')
     returns: 'projects/bar/locations/us-west1/cachedContents/123'
 
   Example:
@@ -72,7 +83,8 @@ def _resource_name(
     client.vertexai = True
     client.project = 'bar'
     client.location = 'us-west1'
-    _resource_name(client, resource_name, collection_identifier='cachedContents')
+    _resource_name(client, resource_name,
+      collection_identifier='cachedContents')
     returns: 'projects/foo/locations/us-central1/cachedContents/123'
 
   Example:
@@ -80,7 +92,8 @@ def _resource_name(
     resource_name = '123'
     # resource_name = 'cachedContents/123'
     client.vertexai = False
-    _resource_name(client, resource_name, collection_identifier='cachedContents')
+    _resource_name(client, resource_name,
+      collection_identifier='cachedContents')
     returns 'cachedContents/123'
 
   Example:
@@ -88,7 +101,8 @@ def _resource_name(
     resource_prefix = 'cachedContents'
     client.vertexai = False
     # client.vertexai = True
-    _resource_name(client, resource_name, collection_identifier='cachedContents')
+    _resource_name(client, resource_name,
+      collection_identifier='cachedContents')
     returns: 'some/wrong/cachedContents/resource/name/123'
 
   Returns:
@@ -99,7 +113,7 @@ def _resource_name(
       # Check if prepending the collection identifier won't violate the
       # collection hierarchy depth.
       and f'{collection_identifier}/{resource_name}'.count('/') + 1
-      == collection_hirearchy_depth
+      == collection_hierarchy_depth
   )
   if client.vertexai:
     if resource_name.startswith('projects/'):
@@ -142,6 +156,7 @@ def t_model(client: _api_client.ApiClient, model: str):
     else:
       return f'models/{model}'
 
+
 def t_models_url(api_client: _api_client.ApiClient, base_models: bool) -> str:
   if api_client.vertexai:
     if base_models:
@@ -155,8 +170,12 @@ def t_models_url(api_client: _api_client.ApiClient, base_models: bool) -> str:
       return 'tunedModels'
 
 
-def t_extract_models(api_client: _api_client.ApiClient, response: dict) -> list[types.Model]:
-  if response.get('models') is not None:
+def t_extract_models(
+    api_client: _api_client.ApiClient, response: dict
+) -> list[types.Model]:
+  if not response:
+    return []
+  elif response.get('models') is not None:
     return response.get('models')
   elif response.get('tunedModels') is not None:
     return response.get('tunedModels')
@@ -181,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:
@@ -194,16 +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(file_uri=part.uri, mime_type=part.mime_type)
   else:
     return part
 
@@ -282,32 +317,234 @@ def t_contents(
     return [t_content(client, contents)]
 
 
-def process_schema(data: dict):
-  if isinstance(data, dict):
-    # Iterate over a copy of keys to allow deletion
-    for key in list(data.keys()):
-      if key == 'title':
-        del data[key]
-      elif key == 'type':
-        data[key] = data[key].upper()
+def handle_null_fields(schema: dict[str, Any]):
+  """Process null fields in the schema so it is compatible with OpenAPI.
+
+  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
+
+  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"
+        }
+      }
+
+    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']
+
+
+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.
+
+  - Removes the `title` field from the schema if the client is not vertexai.
+  - Inlines the $defs.
+
+  Example of a schema before and after (with mldev):
+    Before:
+
+    `schema`
+
+    {
+        'items': {
+            '$ref': '#/$defs/CountryInfo'
+        },
+        'title': 'Placeholder',
+        'type': 'array'
+    }
+
+
+    `defs`
+
+    {
+      'CountryInfo': {
+        'properties': {
+          'continent': {
+              'title': 'Continent',
+              'type': 'string'
+          },
+          'gdp': {
+              'title': 'Gdp',
+              'type': 'integer'}
+          },
+        }
+        'required':['continent', 'gdp'],
+        'title': 'CountryInfo',
+        'type': 'object'
+      }
+    }
+
+    After:
+
+    `schema`
+     {
+        'items': {
+          'properties': {
+            'continent': {
+                'type': 'string'
+            },
+            'gdp': {
+                'type': 'integer'}
+            },
+          }
+          'required':['continent', 'gdp'],
+          'type': 'object'
+        },
+        'type': 'array'
+    }
+  """
+  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
+
+  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:
-        process_schema(data[key])
-  elif isinstance(data, list):
-    for item in data:
-      process_schema(item)
+        ref = defs[ref_key.split('defs/')[-1]]
+        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
 
-  return data
+  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(
-    _: _api_client.ApiClient, origin: Union[types.SchemaDict, Any]
+    client: _api_client.ApiClient, origin: Union[types.SchemaUnionDict, Any]
 ) -> Optional[types.Schema]:
   if not origin:
     return None
   if isinstance(origin, dict):
-    return origin
-  schema = process_schema(origin.model_json_schema())
-  return types.Schema.model_validate(schema)
+    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 = origin.model_dump(exclude_unset=True)
+    process_schema(schema, client)
+    return types.Schema.model_validate(schema)
+
+  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}')
 
 
 def t_speech_config(
@@ -343,10 +580,12 @@ def t_speech_config(
 def t_tool(client: _api_client.ApiClient, origin) -> types.Tool:
   if not origin:
     return None
-  if inspect.isfunction(origin):
+  if inspect.isfunction(origin) or inspect.ismethod(origin):
     return types.Tool(
         function_declarations=[
-            types.FunctionDeclaration.from_function(client, origin)
+            types.FunctionDeclaration.from_callable(
+                client=client, callable=origin
+            )
         ]
     )
   else:
@@ -456,10 +695,25 @@ def t_resolve_operation(api_client: _api_client.ApiClient, struct: dict):
     return struct
 
 
-def t_file_name(api_client: _api_client.ApiClient, name: str):
-  # Remove the files/ prefx since it's added to the url path.
-  if name.startswith('files/'):
-    return name.split('files/')[1]
+def t_file_name(
+    api_client: _api_client.ApiClient, name: Union[str, types.File]
+):
+  # Remove the files/ prefix since it's added to the url path.
+  if isinstance(name, types.File):
+    name = name.name
+
+  if name is None:
+    raise ValueError('File name is required.')
+
+  if name.startswith('https://'):
+    suffix = name.split('files/')[1]
+    match = re.match('[a-z0-9]+', suffix)
+    if match is None:
+      raise ValueError(f'Could not extract file name from URI: {name}')
+    name = match.group(0)
+  elif name.startswith('files/'):
+    name = name.split('files/')[1]
+
   return name
 
 
@@ -481,12 +735,8 @@ def t_tuning_job_status(
 # Some fields don't accept url safe base64 encoding.
 # We shouldn't use this transformer if the backend adhere to Cloud Type
 # format https://cloud.google.com/docs/discovery/type-format.
-# TODO(b/389133914): Remove the hack after Vertex backend fix the issue.
+# TODO(b/389133914,b/390320301): Remove the hack after backend fix the issue.
 def t_bytes(api_client: _api_client.ApiClient, data: bytes) -> str:
   if not isinstance(data, bytes):
     return data
-  if api_client.vertexai:
-    return base64.b64encode(data).decode('ascii')
-  else:
-    return base64.urlsafe_encode(data).decode('ascii')
-
+  return base64.b64encode(data).decode('ascii')