google-genai 0.0.1__py3-none-any.whl

google/genai/_automatic_function_calling_util.py

@@ -0,0 +1,341 @@
+# Copyright 2024 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 inspect
+import types as typing_types
+from typing import Any, Callable, Literal, Union, _GenericAlias, get_args, get_origin
+import pydantic
+from . import types
+
+_py_builtin_type_to_schema_type = {
+    str: 'STRING',
+    int: 'INTEGER',
+    float: 'NUMBER',
+    bool: 'BOOLEAN',
+    list: 'ARRAY',
+    dict: 'OBJECT',
+}
+
+
+def _is_builtin_primitive_or_compound(
+    annotation: inspect.Parameter.annotation,
+) -> bool:
+  return annotation in _py_builtin_type_to_schema_type.keys()
+
+
+def _raise_for_any_of_if_mldev(schema: types.Schema):
+  if schema.any_of:
+    raise ValueError(
+        'AnyOf is not supported in function declaration schema for Google AI.'
+    )
+
+
+def _raise_for_default_if_mldev(schema: types.Schema):
+  if schema.default is not None:
+    raise ValueError(
+        'Default value is not supported in function declaration schema for'
+        ' Google AI.'
+    )
+
+
+def _raise_for_nullable_if_mldev(schema: types.Schema):
+  if schema.nullable:
+    raise ValueError(
+        'Nullable is not supported in function declaration schema for'
+        ' Google AI.'
+    )
+
+
+def _raise_if_schema_unsupported(client, schema: types.Schema):
+  if not client.vertexai:
+    _raise_for_any_of_if_mldev(schema)
+    _raise_for_default_if_mldev(schema)
+    _raise_for_nullable_if_mldev(schema)
+
+
+def _is_default_value_compatible(
+    default_value: Any, annotation: inspect.Parameter.annotation
+) -> bool:
+  # None type is expected to be handled external to this function
+  if _is_builtin_primitive_or_compound(annotation):
+    return isinstance(default_value, annotation)
+
+  if (
+      isinstance(annotation, _GenericAlias)
+      or isinstance(annotation, typing_types.GenericAlias)
+      or isinstance(annotation, typing_types.UnionType)
+  ):
+    origin = get_origin(annotation)
+    if origin in (Union, typing_types.UnionType):
+      return any(
+          _is_default_value_compatible(default_value, arg)
+          for arg in get_args(annotation)
+      )
+
+    if origin is dict:
+      return isinstance(default_value, dict)
+
+    if origin is list:
+      if not isinstance(default_value, list):
+        return False
+      # most tricky case, element in list is union type
+      # need to apply any logic within all
+      # see test case test_generic_alias_complex_array_with_default_value
+      # a: typing.List[int | str | float | bool]
+      # default_value: [1, 'a', 1.1, True]
+      return all(
+          any(
+              _is_default_value_compatible(item, arg)
+              for arg in get_args(annotation)
+          )
+          for item in default_value
+      )
+
+    if origin is Literal:
+      return default_value in get_args(annotation)
+
+  # return False for any other unrecognized annotation
+  # let caller handle the raise
+  return False
+
+
+def _parse_schema_from_parameter(
+    client, param: inspect.Parameter, func_name: str
+) -> types.Schema:
+  """parse schema from parameter.
+
+  from the simplest case to the most complex case.
+  """
+  schema = types.Schema()
+  default_value_error_msg = (
+      f'Default value {param.default} of parameter {param} of function'
+      f' {func_name} is not compatible with the parameter annotation'
+      f' {param.annotation}.'
+  )
+  if _is_builtin_primitive_or_compound(param.annotation):
+    if param.default is not inspect.Parameter.empty:
+      if not _is_default_value_compatible(param.default, param.annotation):
+        raise ValueError(default_value_error_msg)
+      schema.default = param.default
+    schema.type = _py_builtin_type_to_schema_type[param.annotation]
+    _raise_if_schema_unsupported(client, schema)
+    return schema
+  if (
+      isinstance(param.annotation, typing_types.UnionType)
+      # only parse simple UnionType, example int | str | float | bool
+      # complex types.UnionType will be invoked in raise branch
+      and all(
+          (_is_builtin_primitive_or_compound(arg) or arg is type(None))
+          for arg in get_args(param.annotation)
+      )
+  ):
+    schema.type = 'OBJECT'
+    schema.any_of = []
+    unique_types = set()
+    for arg in get_args(param.annotation):
+      if arg.__name__ == 'NoneType':  # Optional type
+        schema.nullable = True
+        continue
+      schema_in_any_of = _parse_schema_from_parameter(
+          client,
+          inspect.Parameter(
+              'item', inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=arg
+          ),
+          func_name,
+      )
+      if (
+          schema_in_any_of.model_dump_json(exclude_none=True)
+          not in unique_types
+      ):
+        schema.any_of.append(schema_in_any_of)
+        unique_types.add(schema_in_any_of.model_dump_json(exclude_none=True))
+    if len(schema.any_of) == 1:  # param: list | None -> Array
+      schema.type = schema.any_of[0].type
+      schema.any_of = None
+    if (
+        param.default is not inspect.Parameter.empty
+        and param.default is not None
+    ):
+      if not _is_default_value_compatible(param.default, param.annotation):
+        raise ValueError(default_value_error_msg)
+      # TODO: b/379715133 - handle pydantic model default value
+      schema.default = param.default
+    _raise_if_schema_unsupported(client, schema)
+    return schema
+  if isinstance(param.annotation, _GenericAlias) or isinstance(
+      param.annotation, typing_types.GenericAlias
+  ):
+    origin = get_origin(param.annotation)
+    args = get_args(param.annotation)
+    if origin is dict:
+      schema.type = 'OBJECT'
+      if param.default is not inspect.Parameter.empty:
+        if not _is_default_value_compatible(param.default, param.annotation):
+          raise ValueError(default_value_error_msg)
+        schema.default = param.default
+      _raise_if_schema_unsupported(client, schema)
+      return schema
+    if origin is Literal:
+      if not all(isinstance(arg, str) for arg in args):
+        raise ValueError(
+            f'Literal type {param.annotation} must be a list of strings.'
+        )
+      schema.type = 'STRING'
+      schema.enum = list(args)
+      if param.default is not inspect.Parameter.empty:
+        if not _is_default_value_compatible(param.default, param.annotation):
+          raise ValueError(default_value_error_msg)
+        schema.default = param.default
+      _raise_if_schema_unsupported(client, schema)
+      return schema
+    if origin is list:
+      schema.type = 'ARRAY'
+      schema.items = _parse_schema_from_parameter(
+          client,
+          inspect.Parameter(
+              'item',
+              inspect.Parameter.POSITIONAL_OR_KEYWORD,
+              annotation=args[0],
+          ),
+          func_name,
+      )
+      if param.default is not inspect.Parameter.empty:
+        if not _is_default_value_compatible(param.default, param.annotation):
+          raise ValueError(default_value_error_msg)
+        schema.default = param.default
+      _raise_if_schema_unsupported(client, schema)
+      return schema
+    if origin is Union:
+      schema.any_of = []
+      schema.type = 'OBJECT'
+      unique_types = set()
+      for arg in args:
+        if arg.__name__ == 'NoneType':  # Optional type
+          schema.nullable = True
+          continue
+        schema_in_any_of = _parse_schema_from_parameter(
+            client,
+            inspect.Parameter(
+                'item',
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=arg,
+            ),
+            func_name,
+        )
+        if (
+            schema_in_any_of.model_dump_json(exclude_none=True)
+            not in unique_types
+        ):
+          schema.any_of.append(schema_in_any_of)
+          unique_types.add(schema_in_any_of.model_dump_json(exclude_none=True))
+      if len(schema.any_of) == 1:  # param: Union[List, None] -> Array
+        schema.type = schema.any_of[0].type
+        schema.any_of = None
+      if (
+          param.default is not None
+          and param.default is not inspect.Parameter.empty
+      ):
+        if not _is_default_value_compatible(param.default, param.annotation):
+          raise ValueError(default_value_error_msg)
+        schema.default = param.default
+      _raise_if_schema_unsupported(client, schema)
+      return schema
+      # all other generic alias will be invoked in raise branch
+  if (
+      inspect.isclass(param.annotation)
+      # for user defined class, we only support pydantic model
+      and issubclass(param.annotation, pydantic.BaseModel)
+  ):
+    if param.default is not inspect.Parameter.empty:
+      # TODO: b/379715133 - handle pydantic model default value
+      raise ValueError(
+          f'Default value {param.default} of Pydantic model{param} of function'
+          f' {func_name} is not supported.'
+      )
+    schema.type = 'OBJECT'
+    schema.properties = {}
+    for field_name, field_info in param.annotation.model_fields.items():
+      schema.properties[field_name] = _parse_schema_from_parameter(
+          client,
+          inspect.Parameter(
+              field_name,
+              inspect.Parameter.POSITIONAL_OR_KEYWORD,
+              annotation=field_info.annotation,
+          ),
+          func_name,
+      )
+    _raise_if_schema_unsupported(client, schema)
+    return schema
+  raise ValueError(
+      f'Failed to parse the parameter {param} of function {func_name} for'
+      ' automatic function calling.Automatic function calling works best with'
+      ' simpler function signature schema,consider manually parse your'
+      f' function declaration for function {func_name}.'
+  )
+
+
+def _get_required_fields(schema: types.Schema) -> list[str]:
+  if not schema.properties:
+    return
+  return [
+      field_name
+      for field_name, field_schema in schema.properties.items()
+      if not field_schema.nullable and field_schema.default is None
+  ]
+
+
+def function_to_declaration(
+    client, func: Callable
+) -> types.FunctionDeclaration:
+  """Converts a function to a FunctionDeclaration."""
+  parameters_properties = {}
+  for name, param in inspect.signature(func).parameters.items():
+    if param.kind in (
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        inspect.Parameter.KEYWORD_ONLY,
+        inspect.Parameter.POSITIONAL_ONLY,
+    ):
+      schema = _parse_schema_from_parameter(client, param, func.__name__)
+      parameters_properties[name] = schema
+  declaration = types.FunctionDeclaration(
+      name=func.__name__,
+      description=func.__doc__,
+  )
+  if parameters_properties:
+    declaration.parameters = types.Schema(
+        type='OBJECT',
+        properties=parameters_properties,
+    )
+    if client.vertexai:
+      declaration.parameters.required = _get_required_fields(
+          declaration.parameters
+      )
+  if not client.vertexai:
+    return declaration
+
+  return_annotation = inspect.signature(func).return_annotation
+  if return_annotation is inspect._empty:
+    return declaration
+
+  declaration.response = _parse_schema_from_parameter(
+      client,
+      inspect.Parameter(
+          'return_value',
+          inspect.Parameter.POSITIONAL_OR_KEYWORD,
+          annotation=return_annotation,
+      ),
+      func.__name__,
+  )
+  return declaration

google/genai/_common.py

@@ -0,0 +1,256 @@
+# Copyright 2024 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.
+#
+
+"""Common utilities for the SDK."""
+
+import base64
+import datetime
+import json
+import typing
+from typing import Union
+import uuid
+
+import pydantic
+from pydantic import alias_generators
+
+from . import _api_client
+
+
+def set_value_by_path(data, keys, value):
+  """Examples:
+
+  set_value_by_path({}, ['a', 'b'], v)
+    -> {'a': {'b': v}}
+  set_value_by_path({}, ['a', 'b[]', c], [v1, v2])
+    -> {'a': {'b': [{'c': v1}, {'c': v2}]}}
+  set_value_by_path({'a': {'b': [{'c': v1}, {'c': v2}]}}, ['a', 'b[]', 'd'], v3)
+    -> {'a': {'b': [{'c': v1, 'd': v3}, {'c': v2, 'd': v3}]}}
+  """
+  if value is None:
+    return
+  for i, key in enumerate(keys[:-1]):
+    if key.endswith('[]'):
+      key_name = key[:-2]
+      if key_name not in data:
+        if isinstance(value, list):
+          data[key_name] = [{} for _ in range(len(value))]
+        else:
+          raise ValueError(
+              f'value {value} must be a list given an array path {key}'
+          )
+      if isinstance(value, list):
+        for j, d in enumerate(data[key_name]):
+          set_value_by_path(d, keys[i + 1 :], value[j])
+      else:
+        for d in data[key_name]:
+          set_value_by_path(d, keys[i + 1 :], value)
+      return
+
+    data = data.setdefault(key, {})
+
+  existing_data = data.get(keys[-1])
+  # If there is an existing value, merge, not overwrite.
+  if existing_data is not None:
+    # Don't overwrite existing non-empty value with new empty value.
+    # This is triggered when handling tuning datasets.
+    if not value:
+      pass
+    # Don't fail when overwriting value with same value
+    elif value == existing_data:
+      pass
+    # Instead of overwriting dictionary with another dictionary, merge them.
+    # This is important for handling training and validation datasets in tuning.
+    elif isinstance(existing_data, dict) and isinstance(value, dict):
+      # Merging dictionaries. Consider deep merging in the future.
+      existing_data.update(value)
+    else:
+      raise ValueError(
+          f'Cannot set value for an existing key. Key: {keys[-1]};'
+          f' Existing value: {existing_data}; New value: {value}.'
+      )
+  else:
+    data[keys[-1]] = value
+
+
+def get_value_by_path(data: object, keys: list[str]):
+  """Examples:
+
+  get_value_by_path({'a': {'b': v}}, ['a', 'b'])
+    -> v
+  get_value_by_path({'a': {'b': [{'c': v1}, {'c': v2}]}}, ['a', 'b[]', 'c'])
+    -> [v1, v2]
+  """
+  if keys == ['_self']:
+    return data
+  for i, key in enumerate(keys):
+    if not data:
+      return None
+    if key.endswith('[]'):
+      key_name = key[:-2]
+      if key_name in data:
+        return [get_value_by_path(d, keys[i + 1 :]) for d in data[key_name]]
+      else:
+        return None
+    else:
+      if key in data:
+        data = data[key]
+      elif isinstance(data, BaseModel) and hasattr(data, key):
+        data = getattr(data, key)
+      else:
+        return None
+  return data
+
+
+class BaseModule:
+
+  def __init__(self, api_client_: _api_client.ApiClient):
+    self.api_client = api_client_
+
+
+def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
+  """Recursively converts a given object to a dictionary.
+
+  If the object is a Pydantic model, it uses the model's `model_dump()` method.
+
+  Args:
+    obj: The object to convert.
+
+  Returns:
+    A dictionary representation of the object.
+  """
+  if isinstance(obj, pydantic.BaseModel):
+    return obj.model_dump(exclude_none=True)
+  elif isinstance(obj, dict):
+    return {key: convert_to_dict(value) for key, value in obj.items()}
+  elif isinstance(obj, list):
+    return [convert_to_dict(item) for item in obj]
+  else:
+    return obj
+
+
+def _remove_extra_fields(
+    model: pydantic.BaseModel, response: dict[str, object]
+) -> None:
+  """Removes extra fields from the response that are not in the model.
+
+  Muates the response in place.
+  """
+
+  key_values = list(response.items())
+
+  for key, value in key_values:
+    # Need to convert to snake case to match model fields names
+    # ex: UsageMetadata
+    alias_map = {
+        field_info.alias: key for key, field_info in model.model_fields.items()
+    }
+
+    if key not in model.model_fields and key not in alias_map:
+      response.pop(key)
+      continue
+
+    key = alias_map.get(key, key)
+
+    annotation = model.model_fields[key].annotation
+
+    # Get the BaseModel if Optional
+    if typing.get_origin(annotation) is Union:
+      annotation = typing.get_args(annotation)[0]
+
+    # if dict, assume BaseModel but also check that field type is not dict
+    # example: FunctionCall.args
+    if isinstance(value, dict) and typing.get_origin(annotation) is not dict:
+      _remove_extra_fields(annotation, value)
+    elif isinstance(value, list):
+      for item in value:
+        # assume a list of dict is list of BaseModel
+        if isinstance(item, dict):
+          _remove_extra_fields(typing.get_args(annotation)[0], item)
+
+
+class BaseModel(pydantic.BaseModel):
+
+  model_config = pydantic.ConfigDict(
+      alias_generator=alias_generators.to_camel,
+      populate_by_name=True,
+      from_attributes=True,
+      protected_namespaces={},
+      extra='forbid',
+      # This allows us to use arbitrary types in the model. E.g. PIL.Image.
+      arbitrary_types_allowed=True,
+  )
+
+  @classmethod
+  def _from_response(
+      cls, response: dict[str, object], kwargs: dict[str, object]
+  ) -> 'BaseModel':
+    # To maintain forward compatibility, we need to remove extra fields from
+    # the response.
+    # We will provide another mechanism to allow users to access these fields.
+    _remove_extra_fields(cls, response)
+    validated_response = cls.model_validate(response)
+    return apply_base64_decoding_for_model(validated_response)
+
+
+def timestamped_unique_name() -> str:
+  """Composes a timestamped unique name.
+
+  Returns:
+      A string representing a unique name.
+  """
+  timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
+  unique_id = uuid.uuid4().hex[0:5]
+  return f'{timestamp}_{unique_id}'
+
+
+def apply_base64_encoding(data: dict[str, object]) -> dict[str, object]:
+  """Applies base64 encoding to bytes values in the given data."""
+  return process_bytes_fields(data, encode=True)
+
+
+def apply_base64_decoding(data: dict[str, object]) -> dict[str, object]:
+  """Applies base64 decoding to bytes values in the given data."""
+  return process_bytes_fields(data, encode=False)
+
+
+def apply_base64_decoding_for_model(data: BaseModel) -> BaseModel:
+  d = data.model_dump(exclude_none=True)
+  d = apply_base64_decoding(d)
+  return data.model_validate(d)
+
+
+def process_bytes_fields(data: dict[str, object], encode=True) -> dict[str, object]:
+  processed_data = {}
+  if not isinstance(data, dict):
+    return data
+  for key, value in data.items():
+    if isinstance(value, bytes):
+      if encode:
+        processed_data[key] = base64.b64encode(value)
+      else:
+        processed_data[key] = base64.b64decode(value)
+    elif isinstance(value, dict):
+      processed_data[key] = process_bytes_fields(value, encode)
+    elif isinstance(value, list):
+      if encode and all(isinstance(v, bytes) for v in value):
+        processed_data[key] = [base64.b64encode(v) for v in value]
+      elif all(isinstance(v, bytes) for v in value):
+        processed_data[key] = [base64.b64decode(v) for v in value]
+      else:
+        processed_data[key] = [process_bytes_fields(v, encode) for v in value]
+    else:
+      processed_data[key] = value
+  return processed_data
+