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

google/genai/_automatic_function_calling_util.py

@@ -14,11 +14,18 @@
 #
 
 import inspect
-import types as typing_types
+import sys
+import types as builtin_types
+import typing
 from typing import Any, Callable, Literal, Union, _GenericAlias, get_args, get_origin
 import pydantic
 from . import types
 
+if sys.version_info >= (3, 10):
+  UnionType = builtin_types.UnionType
+else:
+  UnionType = typing._UnionGenericAlias
+
 _py_builtin_type_to_schema_type = {
     str: 'STRING',
     int: 'INTEGER',
@@ -58,8 +65,8 @@ def _raise_for_nullable_if_mldev(schema: types.Schema):
     )
 
 
-def _raise_if_schema_unsupported(variant: str, schema: types.Schema):
-  if not variant == 'VERTEX_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)
@@ -74,11 +81,11 @@ def _is_default_value_compatible(
 
   if (
       isinstance(annotation, _GenericAlias)
-      or isinstance(annotation, typing_types.GenericAlias)
-      or isinstance(annotation, typing_types.UnionType)
+      or isinstance(annotation, builtin_types.GenericAlias)
+      or isinstance(annotation, UnionType)
   ):
     origin = get_origin(annotation)
-    if origin in (Union, typing_types.UnionType):
+    if origin in (Union, UnionType):
       return any(
           _is_default_value_compatible(default_value, arg)
           for arg in get_args(annotation)
@@ -107,12 +114,13 @@ def _is_default_value_compatible(
       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(
-    variant: str, param: inspect.Parameter, func_name: str
+    client,
+    param: inspect.Parameter,
+    func_name: str,
 ) -> types.Schema:
   """parse schema from parameter.
 
@@ -130,12 +138,12 @@ def _parse_schema_from_parameter(
         raise ValueError(default_value_error_msg)
       schema.default = param.default
     schema.type = _py_builtin_type_to_schema_type[param.annotation]
-    _raise_if_schema_unsupported(variant, schema)
+    _raise_if_schema_unsupported(client, schema)
     return schema
   if (
-      isinstance(param.annotation, typing_types.UnionType)
+      isinstance(param.annotation, UnionType)
       # only parse simple UnionType, example int | str | float | bool
-      # complex types.UnionType will be invoked in raise branch
+      # complex 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)
@@ -149,7 +157,7 @@ def _parse_schema_from_parameter(
         schema.nullable = True
         continue
       schema_in_any_of = _parse_schema_from_parameter(
-          variant,
+          client,
           inspect.Parameter(
               'item', inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=arg
           ),
@@ -171,10 +179,10 @@ def _parse_schema_from_parameter(
       if not _is_default_value_compatible(param.default, param.annotation):
         raise ValueError(default_value_error_msg)
       schema.default = param.default
-    _raise_if_schema_unsupported(variant, schema)
+    _raise_if_schema_unsupported(client, schema)
     return schema
   if isinstance(param.annotation, _GenericAlias) or isinstance(
-      param.annotation, typing_types.GenericAlias
+      param.annotation, builtin_types.GenericAlias
   ):
     origin = get_origin(param.annotation)
     args = get_args(param.annotation)
@@ -184,7 +192,7 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
     if origin is Literal:
       if not all(isinstance(arg, str) for arg in args):
@@ -197,12 +205,12 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
     if origin is list:
       schema.type = 'ARRAY'
       schema.items = _parse_schema_from_parameter(
-          variant,
+          client,
           inspect.Parameter(
               'item',
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -214,7 +222,7 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
     if origin is Union:
       schema.any_of = []
@@ -225,7 +233,7 @@ def _parse_schema_from_parameter(
           schema.nullable = True
           continue
         schema_in_any_of = _parse_schema_from_parameter(
-            variant,
+            client,
             inspect.Parameter(
                 'item',
                 inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -233,6 +241,17 @@ def _parse_schema_from_parameter(
             ),
             func_name,
         )
+        if (
+            len(param.annotation.__args__) == 2
+            and type(None) in param.annotation.__args__
+        ):  # Optional type
+          for optional_arg in param.annotation.__args__:
+            if (
+                hasattr(optional_arg, '__origin__')
+                and optional_arg.__origin__ is list
+            ):
+              # Optional type with list, for example Optional[list[str]]
+              schema.items = schema_in_any_of.items
         if (
             schema_in_any_of.model_dump_json(exclude_none=True)
             not in unique_types
@@ -249,7 +268,7 @@ def _parse_schema_from_parameter(
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
         schema.default = param.default
-      _raise_if_schema_unsupported(variant, schema)
+      _raise_if_schema_unsupported(client, schema)
       return schema
       # all other generic alias will be invoked in raise branch
   if (
@@ -266,7 +285,7 @@ def _parse_schema_from_parameter(
     schema.properties = {}
     for field_name, field_info in param.annotation.model_fields.items():
       schema.properties[field_name] = _parse_schema_from_parameter(
-          variant,
+          client,
           inspect.Parameter(
               field_name,
               inspect.Parameter.POSITIONAL_OR_KEYWORD,
@@ -274,7 +293,9 @@ def _parse_schema_from_parameter(
           ),
           func_name,
       )
-    _raise_if_schema_unsupported(variant, schema)
+    if client.vertexai:
+      schema.required = _get_required_fields(schema)
+    _raise_if_schema_unsupported(client, schema)
     return schema
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'

google/genai/_common.py

@@ -17,6 +17,7 @@
 
 import base64
 import datetime
+import enum
 import typing
 from typing import Union
 import uuid
@@ -112,12 +113,6 @@ def get_value_by_path(data: object, keys: list[str]):
   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.
 
@@ -144,7 +139,7 @@ def _remove_extra_fields(
 ) -> None:
   """Removes extra fields from the response that are not in the model.
 
-  Muates the response in place.
+  Mutates the response in place.
   """
 
   key_values = list(response.items())
@@ -185,7 +180,7 @@ class BaseModel(pydantic.BaseModel):
       alias_generator=alias_generators.to_camel,
       populate_by_name=True,
       from_attributes=True,
-      protected_namespaces={},
+      protected_namespaces=(),
       extra='forbid',
       # This allows us to use arbitrary types in the model. E.g. PIL.Image.
       arbitrary_types_allowed=True,
@@ -208,6 +203,20 @@ class BaseModel(pydantic.BaseModel):
     return self.model_dump(exclude_none=True, mode='json')
 
 
+class CaseInSensitiveEnum(str, enum.Enum):
+  """Case insensitive enum."""
+
+  @classmethod
+  def _missing_(cls, value):
+    try:
+      return cls[value.upper()]  # Try to access directly with uppercase
+    except KeyError:
+      try:
+        return cls[value.lower()]  # Try to access directly with lowercase
+      except KeyError as e:
+        raise ValueError(f"{value} is not a valid {cls.__name__}") from e
+
+
 def timestamped_unique_name() -> str:
   """Composes a timestamped unique name.
 
@@ -219,23 +228,39 @@ def timestamped_unique_name() -> str:
   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."""
+def encode_unserializable_types(data: dict[str, object]) -> dict[str, object]:
+  """Converts unserializable types in dict to json.dumps() compatible types.
+
+  This function is called in models.py after calling convert_to_dict(). The
+  convert_to_dict() can convert pydantic object to dict. However, the input to
+  convert_to_dict() is dict mixed of pydantic object and nested dict(the output
+  of converters). So they may be bytes in the dict and they are out of
+  `ser_json_bytes` control in model_dump(mode='json') called in
+  `convert_to_dict`, as well as datetime deserialization in Pydantic json mode.
+
+  Returns:
+    A dictionary with json.dumps() incompatible type (e.g. bytes datetime)
+    to compatible type (e.g. base64 encoded string, isoformat date string).
+  """
   processed_data = {}
   if not isinstance(data, dict):
     return data
   for key, value in data.items():
     if isinstance(value, bytes):
       processed_data[key] = base64.urlsafe_b64encode(value).decode('ascii')
+    elif isinstance(value, datetime.datetime):
+      processed_data[key] = value.isoformat()
     elif isinstance(value, dict):
-      processed_data[key] = apply_base64_encoding(value)
+      processed_data[key] = encode_unserializable_types(value)
     elif isinstance(value, list):
       if all(isinstance(v, bytes) for v in value):
         processed_data[key] = [
             base64.urlsafe_b64encode(v).decode('ascii') for v in value
         ]
+      if all(isinstance(v, datetime.datetime) for v in value):
+        processed_data[key] = [v.isoformat() for v in value]
       else:
-        processed_data[key] = [apply_base64_encoding(v) for v in value]
+        processed_data[key] = [encode_unserializable_types(v) for v in value]
     else:
       processed_data[key] = value
   return processed_data

google/genai/_extra_utils.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 #
 
-"""Extra utils depending on types that are shared between sync and async modules.
-"""
+"""Extra utils depending on types that are shared between sync and async modules."""
 
 import inspect
 import logging
-from typing import Any, Callable, Dict, get_args, get_origin, Optional, types as typing_types, Union
+import typing
+from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
+import sys
 
 import pydantic
 
@@ -26,6 +27,10 @@ from . import _common
 from . import errors
 from . import types
 
+if sys.version_info >= (3, 10):
+  from types import UnionType
+else:
+  UnionType = typing._UnionGenericAlias
 
 _DEFAULT_MAX_REMOTE_CALLS_AFC = 10
 
@@ -78,8 +83,8 @@ def get_function_map(
         if inspect.iscoroutinefunction(tool):
           raise errors.UnsupportedFunctionError(
               f'Function {tool.__name__} is a coroutine function, which is not'
-              ' supported for automatic function calling. Please manually invoke'
-              f' {tool.__name__} to get the function response.'
+              ' supported for automatic function calling. Please manually'
+              f' invoke {tool.__name__} to get the function response.'
           )
         function_map[tool.__name__] = tool
   return function_map
@@ -116,7 +121,7 @@ def convert_if_exist_pydantic_model(
     try:
       return annotation(**value)
     except pydantic.ValidationError as e:
-      raise errors.UnkownFunctionCallArgumentError(
+      raise errors.UnknownFunctionCallArgumentError(
           f'Failed to parse parameter {param_name} for function'
           f' {func_name} from function call part because function call argument'
           f' value {value} is not compatible with parameter annotation'
@@ -135,11 +140,13 @@ def convert_if_exist_pydantic_model(
         for k, v in value.items()
     }
   # example 1: typing.Union[int, float]
-  # example 2: int | float equivalent to typing.types.UnionType[int, float]
-  if get_origin(annotation) in (Union, typing_types.UnionType):
+  # example 2: int | float equivalent to UnionType[int, float]
+  if get_origin(annotation) in (Union, UnionType):
     for arg in get_args(annotation):
-      if isinstance(value, arg) or (
-          isinstance(value, dict) and _is_annotation_pydantic_model(arg)
+      if (
+          (get_args(arg) and get_origin(arg) is list)
+          or isinstance(value, arg)
+          or (isinstance(value, dict) and _is_annotation_pydantic_model(arg))
       ):
         try:
           return convert_if_exist_pydantic_model(
@@ -150,7 +157,7 @@ def convert_if_exist_pydantic_model(
         except pydantic.ValidationError:
           continue
     # if none of the union type is matched, raise error
-    raise errors.UnkownFunctionCallArgumentError(
+    raise errors.UnknownFunctionCallArgumentError(
         f'Failed to parse parameter {param_name} for function'
         f' {func_name} from function call part because function call argument'
         f' value {value} cannot be converted to parameter annotation'
@@ -161,7 +168,7 @@ def convert_if_exist_pydantic_model(
   if isinstance(value, int) and annotation is float:
     return value
   if not isinstance(value, annotation):
-    raise errors.UnkownFunctionCallArgumentError(
+    raise errors.UnknownFunctionCallArgumentError(
         f'Failed to parse parameter {param_name} for function {func_name} from'
         f' function call part because function call argument value {value} is'
         f' not compatible with parameter annotation {annotation}.'
@@ -209,7 +216,9 @@ def get_function_response_parts(
       response = {'result': invoke_function_from_dict_args(args, func)}
     except Exception as e:  # pylint: disable=broad-except
       response = {'error': str(e)}
-    func_response = types.Part.from_function_response(func_name, response)
+    func_response = types.Part.from_function_response(
+        name=func_name, response=response
+    )
 
     func_response_parts.append(func_response)
   return func_response_parts
@@ -231,8 +240,7 @@ def should_disable_afc(
       and config_model.automatic_function_calling
       and config_model.automatic_function_calling.maximum_remote_calls
       is not None
-      and int(config_model.automatic_function_calling.maximum_remote_calls)
-      <= 0
+      and int(config_model.automatic_function_calling.maximum_remote_calls) <= 0
   ):
     logging.warning(
         'max_remote_calls in automatic_function_calling_config'
@@ -294,6 +302,7 @@ def get_max_remote_calls_afc(
     return _DEFAULT_MAX_REMOTE_CALLS_AFC
   return int(config_model.automatic_function_calling.maximum_remote_calls)
 
+
 def should_append_afc_history(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> bool:
@@ -302,9 +311,6 @@ def should_append_afc_history(
       if config and isinstance(config, dict)
       else config
   )
-  if (
-      not config_model
-      or not config_model.automatic_function_calling
-  ):
+  if not config_model or not config_model.automatic_function_calling:
     return True
   return not config_model.automatic_function_calling.ignore_call_history

google/genai/_replay_api_client.py

@@ -17,11 +17,12 @@
 
 import base64
 import copy
+import datetime
 import inspect
+import io
 import json
 import os
 import re
-import datetime
 from typing import Any, Literal, Optional, Union
 
 import google.auth
@@ -32,9 +33,9 @@ from ._api_client import ApiClient
 from ._api_client import HttpOptions
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
-from ._api_client import RequestJsonEncoder
 from ._common import BaseModel
 
+
 def _redact_version_numbers(version_string: str) -> str:
   """Redacts version numbers in the form x.y.z from a string."""
   return re.sub(r'\d+\.\d+\.\d+', '{VERSION_NUMBER}', version_string)
@@ -145,6 +146,7 @@ class ReplayResponse(BaseModel):
   status_code: int = 200
   headers: dict[str, str]
   body_segments: list[dict[str, object]]
+  byte_segments: Optional[list[bytes]] = None
   sdk_response_segments: list[dict[str, object]]
 
   def model_post_init(self, __context: Any) -> None:
@@ -264,17 +266,13 @@ class ReplayApiClient(ApiClient):
     replay_file_path = self._get_replay_file_path()
     os.makedirs(os.path.dirname(replay_file_path), exist_ok=True)
     with open(replay_file_path, 'w') as f:
-      f.write(
-          json.dumps(
-              self.replay_session.model_dump(mode='json'), indent=2, cls=ResponseJsonEncoder
-          )
-      )
+      f.write(self.replay_session.model_dump_json(exclude_unset=True, indent=2))
     self.replay_session = None
 
   def _record_interaction(
       self,
       http_request: HttpRequest,
-      http_response: Union[HttpResponse, errors.APIError],
+      http_response: Union[HttpResponse, errors.APIError, bytes],
   ):
     if not self._should_update_replay():
       return
@@ -289,6 +287,9 @@ class ReplayApiClient(ApiClient):
       response = ReplayResponse(
           headers=dict(http_response.headers),
           body_segments=list(http_response.segments()),
+          byte_segments=[
+              seg[:100] + b'...' for seg in http_response.byte_segments()
+          ],
           status_code=http_response.status_code,
           sdk_response_segments=[],
       )
@@ -322,11 +323,7 @@ class ReplayApiClient(ApiClient):
     # so that the comparison is fair.
     _redact_request_body(request_data_copy)
 
-    # Need to call dumps() and loads() to convert dict bytes values to strings.
-    # Because the expected_request_body dict never contains bytes values.
-    actual_request_body = [
-        json.loads(json.dumps(request_data_copy, cls=RequestJsonEncoder))
-    ]
+    actual_request_body = [request_data_copy]
     expected_request_body = interaction.request.body_segments
     assert actual_request_body == expected_request_body, (
         'Request body mismatch:\n'
@@ -349,6 +346,7 @@ class ReplayApiClient(ApiClient):
             json.dumps(segment)
             for segment in interaction.response.body_segments
         ],
+        byte_stream=interaction.response.byte_segments,
     )
 
   def _verify_response(self, response_model: BaseModel):
@@ -368,7 +366,9 @@ class ReplayApiClient(ApiClient):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
     actual = response_model.model_dump(exclude_none=True, mode='json')
-    expected = interaction.response.sdk_response_segments[self._sdk_response_index]
+    expected = interaction.response.sdk_response_segments[
+        self._sdk_response_index
+    ]
     assert (
         actual == expected
     ), f'SDK response mismatch:\nActual: {actual}\nExpected: {expected}'
@@ -397,15 +397,26 @@ class ReplayApiClient(ApiClient):
         # segments since the stream has been consumed.
       else:
         self._record_interaction(http_request, result)
-      _debug_print('api mode result: %s' % result.text)
+      _debug_print('api mode result: %s' % result.json)
       return result
     else:
       return self._build_response_from_replay(http_request)
 
-  def upload_file(self, file_path: str, upload_url: str, upload_size: int):
-    request = HttpRequest(
-        method='POST', url='', data={'file_path': file_path}, headers={}
-    )
+  def upload_file(self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int):
+    if isinstance(file_path, io.IOBase):
+      offset = file_path.tell()
+      content = file_path.read()
+      file_path.seek(offset, os.SEEK_SET)
+      request = HttpRequest(
+          method='POST',
+          url='',
+          data={'bytes': base64.b64encode(content).decode('utf-8')},
+          headers={}
+      )
+    else:
+      request = HttpRequest(
+          method='POST', url='', data={'file_path': file_path}, headers={}
+      )
     if self._should_call_api():
       try:
         result = super().upload_file(file_path, upload_url, upload_size)
@@ -418,20 +429,21 @@ class ReplayApiClient(ApiClient):
       self._record_interaction(request, HttpResponse({}, [json.dumps(result)]))
       return result
     else:
-      return self._build_response_from_replay(request).text
-
-
-# TODO(b/389693448): Cleanup datetime hacks.
-class ResponseJsonEncoder(json.JSONEncoder):
-  """The replay test json encoder for response.
-  """
-  def default(self, o):
-    if isinstance(o, datetime.datetime):
-      # dt.isoformat() prints "2024-11-15T23:27:45.624657+00:00"
-      # but replay files want "2024-11-15T23:27:45.624657Z"
-      if o.isoformat().endswith('+00:00'):
-        return o.isoformat().replace('+00:00', 'Z')
-      else:
-        return o.isoformat()
+      return self._build_response_from_replay(request).json
+
+  def _download_file_request(self, request):
+    self._initialize_replay_session_if_not_loaded()
+    if self._should_call_api():
+      try:
+        result = super()._download_file_request(request)
+      except HTTPError as e:
+        result = HttpResponse(
+            e.response.headers, [json.dumps({'reason': e.response.reason})]
+        )
+        result.status_code = e.response.status_code
+        raise e
+      self._record_interaction(request, result)
+      return result
     else:
-      return super().default(o)
+      return self._build_response_from_replay(request)
+

google/genai/_test_api_client.py

@@ -132,7 +132,7 @@ async def test_async_request_streamed_non_blocking(
 
   chunks = []
   start_time = time.time()
-  async for chunk in api_client.async_request_streamed(
+  async for chunk in await api_client.async_request_streamed(
       http_method, path, request_dict
   ):
     chunks.append(chunk)