google-genai 1.4.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

google/genai/_extra_utils.py

@@ -17,9 +17,9 @@
 
 import inspect
 import logging
+import sys
 import typing
 from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
-import sys
 
 import pydantic
 
@@ -37,6 +37,15 @@ _DEFAULT_MAX_REMOTE_CALLS_AFC = 10
 logger = logging.getLogger('google_genai.models')
 
 
+def _create_generate_content_config_model(
+    config: types.GenerateContentConfigOrDict,
+) -> types.GenerateContentConfig:
+  if isinstance(config, dict):
+    return types.GenerateContentConfig(**config)
+  else:
+    return config
+
+
 def format_destination(
     src: str,
     config: Optional[types.CreateBatchJobConfigOrDict] = None,
@@ -69,16 +78,12 @@ def format_destination(
 
 def get_function_map(
     config: Optional[types.GenerateContentConfigOrDict] = None,
-) -> dict[str, object]:
+) -> dict[str, Callable]:
   """Returns a function map from the config."""
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
-  function_map: dict[str, object] = {}
-  if not config_model:
+  function_map: dict[str, Callable] = {}
+  if not config:
     return function_map
+  config_model = _create_generate_content_config_model(config)
   if config_model.tools:
     for tool in config_model.tools:
       if callable(tool):
@@ -92,6 +97,16 @@ def get_function_map(
   return function_map
 
 
+def convert_number_values_for_dict_function_call_args(
+    args: dict[str, Any],
+) -> dict[str, Any]:
+  """Converts float values in dict with no decimal to integers."""
+  return {
+      key: convert_number_values_for_function_call_args(value)
+      for key, value in args.items()
+  }
+
+
 def convert_number_values_for_function_call_args(
     args: Union[dict[str, object], list[object], object],
 ) -> Union[dict[str, object], list[object], object]:
@@ -210,26 +225,35 @@ def invoke_function_from_dict_args(
 
 def get_function_response_parts(
     response: types.GenerateContentResponse,
-    function_map: dict[str, object],
+    function_map: dict[str, Callable],
 ) -> list[types.Part]:
   """Returns the function response parts from the response."""
   func_response_parts = []
-  for part in response.candidates[0].content.parts:
-    if not part.function_call:
-      continue
-    func_name = part.function_call.name
-    func = function_map[func_name]
-    args = convert_number_values_for_function_call_args(part.function_call.args)
-    func_response: dict[str, Any]
-    try:
-      func_response = {'result': invoke_function_from_dict_args(args, func)}
-    except Exception as e:  # pylint: disable=broad-except
-      func_response = {'error': str(e)}
-    func_response_part = types.Part.from_function_response(
-        name=func_name, response=func_response
-    )
-
-    func_response_parts.append(func_response_part)
+  if (
+      response.candidates is not None
+      and isinstance(response.candidates[0].content, types.Content)
+      and response.candidates[0].content.parts is not None
+  ):
+    for part in response.candidates[0].content.parts:
+      if not part.function_call:
+        continue
+      func_name = part.function_call.name
+      if func_name is not None and part.function_call.args is not None:
+        func = function_map[func_name]
+        args = convert_number_values_for_dict_function_call_args(
+            part.function_call.args
+        )
+        func_response: dict[str, Any]
+        try:
+          func_response = {
+              'result': invoke_function_from_dict_args(args, func)
+          }
+        except Exception as e:  # pylint: disable=broad-except
+          func_response = {'error': str(e)}
+        func_response_part = types.Part.from_function_response(
+            name=func_name, response=func_response
+        )
+        func_response_parts.append(func_response_part)
   return func_response_parts
 
 
@@ -237,12 +261,9 @@ def should_disable_afc(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> bool:
   """Returns whether automatic function calling is enabled."""
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
-
+  if not config:
+    return False
+  config_model = _create_generate_content_config_model(config)
   # If max_remote_calls is less or equal to 0, warn and disable AFC.
   if (
       config_model
@@ -261,8 +282,7 @@ def should_disable_afc(
 
   # Default to enable AFC if not specified.
   if (
-      not config_model
-      or not config_model.automatic_function_calling
+      not config_model.automatic_function_calling
       or config_model.automatic_function_calling.disable is None
   ):
     return False
@@ -295,20 +315,17 @@ def should_disable_afc(
 def get_max_remote_calls_afc(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> int:
+  if not config:
+    return _DEFAULT_MAX_REMOTE_CALLS_AFC
   """Returns the remaining remote calls for automatic function calling."""
   if should_disable_afc(config):
     raise ValueError(
         'automatic function calling is not enabled, but SDK is trying to get'
         ' max remote calls.'
     )
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
+  config_model = _create_generate_content_config_model(config)
   if (
-      not config_model
-      or not config_model.automatic_function_calling
+      not config_model.automatic_function_calling
       or config_model.automatic_function_calling.maximum_remote_calls is None
   ):
     return _DEFAULT_MAX_REMOTE_CALLS_AFC
@@ -318,11 +335,9 @@ def get_max_remote_calls_afc(
 def should_append_afc_history(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> bool:
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
-  if not config_model or not config_model.automatic_function_calling:
+  if not config:
+    return True
+  config_model = _create_generate_content_config_model(config)
+  if not config_model.automatic_function_calling:
     return True
   return not config_model.automatic_function_calling.ignore_call_history

google/genai/_replay_api_client.py

@@ -109,7 +109,8 @@ def _redact_project_location_path(path: str) -> str:
     return path
 
 
-def _redact_request_body(body: dict[str, object]) -> dict[str, object]:
+def _redact_request_body(body: dict[str, object]):
+  """Redacts fields in the request body in place."""
   for key, value in body.items():
     if isinstance(value, str):
       body[key] = _redact_project_location_path(value)
@@ -302,13 +303,24 @@ class ReplayApiClient(BaseApiClient):
           status_code=http_response.status_code,
           sdk_response_segments=[],
       )
-    else:
+    elif isinstance(http_response, errors.APIError):
       response = ReplayResponse(
           headers=dict(http_response.response.headers),
           body_segments=[http_response._to_replay_record()],
           status_code=http_response.code,
           sdk_response_segments=[],
       )
+    elif isinstance(http_response, bytes):
+      response = ReplayResponse(
+          headers={},
+          body_segments=[],
+          byte_segments=[http_response],
+          sdk_response_segments=[],
+      )
+    else:
+      raise ValueError(
+          'Unsupported http_response type: ' + str(type(http_response))
+      )
     self.replay_session.interactions.append(
         ReplayInteraction(request=request, response=response)
     )
@@ -471,6 +483,43 @@ class ReplayApiClient(BaseApiClient):
     else:
       return self._build_response_from_replay(request).json
 
+  async def async_upload_file(
+      self,
+      file_path: Union[str, io.IOBase],
+      upload_url: str,
+      upload_size: int,
+  ) -> str:
+    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():
+      result: Union[str, HttpResponse]
+      try:
+        result = await super().async_upload_file(
+            file_path, upload_url, upload_size
+        )
+      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, HttpResponse({}, [json.dumps(result)]))
+      return result
+    else:
+      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():
@@ -486,3 +535,22 @@ class ReplayApiClient(BaseApiClient):
       return result
     else:
       return self._build_response_from_replay(request)
+
+  async def async_download_file(self, path: str, http_options):
+    self._initialize_replay_session_if_not_loaded()
+    request = self._build_request(
+        'get', path=path, request_dict={}, http_options=http_options
+    )
+    if self._should_call_api():
+      try:
+        result = await super().async_download_file(path, http_options)
+      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 self._build_response_from_replay(request).byte_stream[0]

google/genai/_transformers.py

@@ -26,9 +26,7 @@ import sys
 import time
 import types as builtin_types
 import typing
-from typing import Any, GenericAlias, Optional, Union
-
-import types as builtin_types
+from typing import Any, GenericAlias, Optional, Union  # type: ignore[attr-defined]
 
 if typing.TYPE_CHECKING:
   import PIL.Image
@@ -43,10 +41,11 @@ logger = logging.getLogger('google_genai._transformers')
 if sys.version_info >= (3, 10):
   VersionedUnionType = builtin_types.UnionType
   _UNION_TYPES = (typing.Union, builtin_types.UnionType)
+  from typing import TypeGuard
 else:
   VersionedUnionType = typing._UnionGenericAlias
   _UNION_TYPES = (typing.Union,)
-
+  from typing_extensions import TypeGuard
 
 def _resource_name(
     client: _api_client.BaseApiClient,
@@ -165,7 +164,9 @@ def t_model(client: _api_client.BaseApiClient, model: str):
       return f'models/{model}'
 
 
-def t_models_url(api_client: _api_client.BaseApiClient, base_models: bool) -> str:
+def t_models_url(
+    api_client: _api_client.BaseApiClient, base_models: bool
+) -> str:
   if api_client.vertexai:
     if base_models:
       return 'publishers/google/models'
@@ -179,8 +180,9 @@ def t_models_url(api_client: _api_client.BaseApiClient, base_models: bool) -> st
 
 
 def t_extract_models(
-    api_client: _api_client.BaseApiClient, response: dict
-) -> list[types.Model]:
+    api_client: _api_client.BaseApiClient,
+    response: dict[str, list[types.ModelDict]],
+) -> Optional[list[types.ModelDict]]:
   if not response:
     return []
   elif response.get('models') is not None:
@@ -240,9 +242,7 @@ def pil_to_blob(img) -> types.Blob:
   return types.Blob(mime_type=mime_type, data=data)
 
 
-def t_part(
-    client: _api_client.BaseApiClient, part: Optional[types.PartUnionDict]
-) -> types.Part:
+def t_part(part: Optional[types.PartUnionDict]) -> types.Part:
   try:
     import PIL.Image
 
@@ -268,22 +268,21 @@ def t_part(
 
 
 def t_parts(
-    client: _api_client.BaseApiClient,
-    parts: Optional[Union[list[types.PartUnionDict], types.PartUnionDict]],
+    parts: Optional[Union[list[types.PartUnionDict], types.PartUnionDict, list[types.Part]]],
 ) -> list[types.Part]:
   #
   if parts is None or (isinstance(parts, list) and not parts):
     raise ValueError('content parts are required.')
   if isinstance(parts, list):
-    return [t_part(client, part) for part in parts]
+    return [t_part(part) for part in parts]
   else:
-    return [t_part(client, parts)]
+    return [t_part(parts)]
 
 
 def t_image_predictions(
     client: _api_client.BaseApiClient,
     predictions: Optional[Iterable[Mapping[str, Any]]],
-) -> list[types.GeneratedImage]:
+) -> Optional[list[types.GeneratedImage]]:
   if not predictions:
     return None
   images = []
@@ -333,22 +332,35 @@ def t_content(
 def t_contents_for_embed(
     client: _api_client.BaseApiClient,
     contents: Union[list[types.Content], list[types.ContentDict], ContentType],
-):
-  if client.vertexai and isinstance(contents, list):
-    # TODO: Assert that only text is supported.
-    return [t_content(client, content).parts[0].text for content in contents]
-  elif client.vertexai:
-    return [t_content(client, contents).parts[0].text]
-  elif isinstance(contents, list):
-    return [t_content(client, content) for content in contents]
+) -> Union[list[str], list[types.Content]]:
+  if isinstance(contents, list):
+    transformed_contents = [t_content(client, content) for content in contents]
   else:
-    return [t_content(client, contents)]
+    transformed_contents = [t_content(client, contents)]
+
+  if client.vertexai:
+    text_parts = []
+    for content in transformed_contents:
+      if content is not None:
+        if isinstance(content, dict):
+          content = types.Content.model_validate(content)
+        if content.parts is not None:
+          for part in content.parts:
+            if part.text:
+              text_parts.append(part.text)
+            else:
+              logger.warning(
+                  f'Non-text part found, only returning text parts.'
+              )
+    return text_parts
+  else:
+    return transformed_contents
 
 
 def t_contents(
     client: _api_client.BaseApiClient,
     contents: Optional[
-        Union[types.ContentListUnion, types.ContentListUnionDict]
+        Union[types.ContentListUnion, types.ContentListUnionDict, types.Content]
     ],
 ) -> list[types.Content]:
   if contents is None or (isinstance(contents, list) and not contents):
@@ -366,7 +378,7 @@ def t_contents(
   result: list[types.Content] = []
   accumulated_parts: list[types.Part] = []
 
-  def _is_part(part: types.PartUnionDict) -> bool:
+  def _is_part(part: Union[types.PartUnionDict, Any]) -> TypeGuard[types.PartUnionDict]:
     if (
         isinstance(part, str)
         or isinstance(part, types.File)
@@ -408,7 +420,7 @@ def t_contents(
       accumulated_parts: list[types.Part],
       current_part: types.PartUnionDict,
   ):
-    current_part = t_part(client, current_part)
+    current_part = t_part(current_part)
     if _is_user_part(current_part) == _are_user_parts(accumulated_parts):
       accumulated_parts.append(current_part)
     else:
@@ -430,11 +442,11 @@ def t_contents(
     ):
       _append_accumulated_parts_as_content(result, accumulated_parts)
       if isinstance(content, list):
-        result.append(types.UserContent(parts=content))
+        result.append(types.UserContent(parts=content))  # type: ignore[arg-type]
       else:
         result.append(content)
-    elif (_is_part(content)): # type: ignore
-      _handle_current_part(result, accumulated_parts, content) # type: ignore
+    elif (_is_part(content)):
+      _handle_current_part(result, accumulated_parts, content)
     elif isinstance(content, dict):
       # PactDict is already handled in _is_part
       result.append(types.Content.model_validate(content))
@@ -500,14 +512,14 @@ def handle_null_fields(schema: dict[str, Any]):
         schema['anyOf'].remove({'type': 'null'})
         if len(schema['anyOf']) == 1:
           # If there is only one type left after removing null, remove the anyOf field.
-          for key,val in schema['anyOf'][0].items():
+          for key, val in schema['anyOf'][0].items():
             schema[key] = val
           del schema['anyOf']
 
 
 def process_schema(
     schema: dict[str, Any],
-    client: Optional[_api_client.BaseApiClient] = None,
+    client: _api_client.BaseApiClient,
     defs: Optional[dict[str, Any]] = None,
     *,
     order_properties: bool = True,
@@ -570,12 +582,13 @@ def process_schema(
         'type': 'array'
     }
   """
-  if client and not client.vertexai:
+  if not client.vertexai:
     schema.pop('title', None)
 
     if schema.get('default') is not None:
       raise ValueError(
-          'Default value is not supported in the response schema for the Gemini API.'
+          'Default value is not supported in the response schema for the Gemini'
+          ' API.'
       )
 
   if schema.get('title') == 'PlaceholderLiteralEnum':
@@ -596,18 +609,15 @@ def process_schema(
   # After removing null fields, Optional fields with only one possible type
   # will have a $ref key that needs to be flattened
   # For example: {'default': None, 'description': 'Name of the person', 'nullable': True, '$ref': '#/$defs/TestPerson'}
-  if schema.get('$ref', None):
-    ref = defs[schema.get('$ref').split('defs/')[-1]]
+  schema_ref = schema.get('$ref', None)
+  if schema_ref is not None:
+    ref = defs[schema_ref.split('defs/')[-1]]
     for schema_key in list(ref.keys()):
       schema[schema_key] = ref[schema_key]
     del schema['$ref']
 
   any_of = schema.get('anyOf', None)
   if any_of is not None:
-    if not client.vertexai:
-      raise ValueError(
-          'AnyOf is not supported in the response schema for the Gemini API.'
-      )
     for sub_schema in any_of:
       # $ref is present in any_of if the schema is a union of Pydantic classes
       ref_key = sub_schema.get('$ref', None)
@@ -670,7 +680,7 @@ def process_schema(
 
 
 def _process_enum(
-    enum: EnumMeta, client: Optional[_api_client.BaseApiClient] = None  # type: ignore
+    enum: EnumMeta, client: _api_client.BaseApiClient
 ) -> types.Schema:
   for member in enum:  # type: ignore
     if not isinstance(member.value, str):
@@ -680,7 +690,7 @@ def _process_enum(
       )
 
   class Placeholder(pydantic.BaseModel):
-    placeholder: enum
+    placeholder: enum  # type: ignore[valid-type]
 
   enum_schema = Placeholder.model_json_schema()
   process_schema(enum_schema, client)
@@ -688,12 +698,19 @@ def _process_enum(
   return types.Schema.model_validate(enum_schema)
 
 
+def _is_type_dict_str_any(origin: Union[types.SchemaUnionDict, Any]) -> TypeGuard[dict[str, Any]]:
+  """Verifies the schema is of type dict[str, Any] for mypy type checking."""
+  return isinstance(origin, dict) and all(
+      isinstance(key, str) for key in origin
+  )
+
+
 def t_schema(
     client: _api_client.BaseApiClient, origin: Union[types.SchemaUnionDict, Any]
 ) -> Optional[types.Schema]:
   if not origin:
     return None
-  if isinstance(origin, dict):
+  if isinstance(origin, dict) and _is_type_dict_str_any(origin):
     process_schema(origin, client, order_properties=False)
     return types.Schema.model_validate(origin)
   if isinstance(origin, EnumMeta):
@@ -724,7 +741,7 @@ def t_schema(
   ):
 
     class Placeholder(pydantic.BaseModel):
-      placeholder: origin
+      placeholder: origin  # type: ignore[valid-type]
 
     schema = Placeholder.model_json_schema()
     process_schema(schema, client)
@@ -735,7 +752,8 @@ def t_schema(
 
 
 def t_speech_config(
-    _: _api_client.BaseApiClient, origin: Union[types.SpeechConfigUnionDict, Any]
+    _: _api_client.BaseApiClient,
+    origin: Union[types.SpeechConfigUnionDict, Any],
 ) -> Optional[types.SpeechConfig]:
   if not origin:
     return None
@@ -750,7 +768,10 @@ def t_speech_config(
   if (
       isinstance(origin, dict)
       and 'voice_config' in origin
+      and origin['voice_config'] is not None
       and 'prebuilt_voice_config' in origin['voice_config']
+      and origin['voice_config']['prebuilt_voice_config'] is not None
+      and 'voice_name' in origin['voice_config']['prebuilt_voice_config']
   ):
     return types.SpeechConfig(
         voice_config=types.VoiceConfig(
@@ -764,7 +785,7 @@ def t_speech_config(
   raise ValueError(f'Unsupported speechConfig type: {type(origin)}')
 
 
-def t_tool(client: _api_client.BaseApiClient, origin) -> types.Tool:
+def t_tool(client: _api_client.BaseApiClient, origin) -> Optional[types.Tool]:
   if not origin:
     return None
   if inspect.isfunction(origin) or inspect.ismethod(origin):
@@ -790,12 +811,16 @@ def t_tools(
   for tool in origin:
     transformed_tool = t_tool(client, tool)
     # All functions should be merged into one tool.
-    if transformed_tool.function_declarations:
-      function_tool.function_declarations += (
+    if transformed_tool is not None:
+      if (
           transformed_tool.function_declarations
-      )
-    else:
-      tools.append(transformed_tool)
+          and function_tool.function_declarations is not None
+      ):
+        function_tool.function_declarations += (
+            transformed_tool.function_declarations
+        )
+      else:
+        tools.append(transformed_tool)
   if function_tool.function_declarations:
     tools.append(function_tool)
   return tools
@@ -883,15 +908,28 @@ def t_resolve_operation(api_client: _api_client.BaseApiClient, struct: dict):
 
 
 def t_file_name(
-    api_client: _api_client.BaseApiClient, name: Optional[Union[str, types.File]]
+    api_client: _api_client.BaseApiClient,
+    name: Optional[Union[str, types.File, types.Video, types.GeneratedVideo]],
 ):
   # Remove the files/ prefix since it's added to the url path.
   if isinstance(name, types.File):
     name = name.name
+  elif isinstance(name, types.Video):
+    name = name.uri
+  elif isinstance(name, types.GeneratedVideo):
+    if name.video is not None:
+      name = name.video.uri
+    else:
+      name = None
 
   if name is None:
     raise ValueError('File name is required.')
 
+  if not isinstance(name, str):
+    raise ValueError(
+        f'Could not convert object of type `{type(name)}` to a file name.'
+    )
+
   if name.startswith('https://'):
     suffix = name.split('files/')[1]
     match = re.match('[a-z0-9]+', suffix)
@@ -906,16 +944,19 @@ def t_file_name(
 
 def t_tuning_job_status(
     api_client: _api_client.BaseApiClient, status: str
-) -> types.JobState:
+) -> Union[types.JobState, str]:
   if status == 'STATE_UNSPECIFIED':
-    return 'JOB_STATE_UNSPECIFIED'
+    return types.JobState.JOB_STATE_UNSPECIFIED
   elif status == 'CREATING':
-    return 'JOB_STATE_RUNNING'
+    return types.JobState.JOB_STATE_RUNNING
   elif status == 'ACTIVE':
-    return 'JOB_STATE_SUCCEEDED'
+    return types.JobState.JOB_STATE_SUCCEEDED
   elif status == 'FAILED':
-    return 'JOB_STATE_FAILED'
+    return types.JobState.JOB_STATE_FAILED
   else:
+    for state in types.JobState:
+      if str(state.value) == status:
+        return state
     return status