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

google/genai/_api_client.py

@@ -19,6 +19,7 @@
 The BaseApiClient is intended to be a private module and is subject to change.
 """
 
+import anyio
 import asyncio
 import copy
 from dataclasses import dataclass
@@ -44,6 +45,7 @@ from . import version
 from .types import HttpOptions, HttpOptionsDict, HttpOptionsOrDict
 
 logger = logging.getLogger('google_genai._api_client')
+CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunk size
 
 
 def _append_library_version_headers(headers: dict[str, str]) -> None:
@@ -118,8 +120,9 @@ def _load_auth(*, project: Union[str, None]) -> tuple[Credentials, str]:
   return credentials, project
 
 
-def _refresh_auth(credentials: Credentials) -> None:
+def _refresh_auth(credentials: Credentials) -> Credentials:
   credentials.refresh(Request())
+  return credentials
 
 
 @dataclass
@@ -147,7 +150,7 @@ class HttpResponse:
 
   def __init__(
       self,
-      headers: dict[str, str],
+      headers: Union[dict[str, str], httpx.Headers],
       response_stream: Union[Any, str] = None,
       byte_stream: Union[Any, bytes] = None,
   ):
@@ -200,14 +203,19 @@ class HttpResponse:
         yield c
     else:
       # Iterator of objects retrieved from the API.
-      async for chunk in self.response_stream.aiter_lines():
-        # This is httpx.Response.
-        if chunk:
-          # In async streaming mode, the chunk of JSON is prefixed with "data:"
-          # which we must strip before parsing.
-          if chunk.startswith('data: '):
-            chunk = chunk[len('data: ') :]
-          yield json.loads(chunk)
+      if hasattr(self.response_stream, 'aiter_lines'):
+        async for chunk in self.response_stream.aiter_lines():
+          # This is httpx.Response.
+          if chunk:
+            # In async streaming mode, the chunk of JSON is prefixed with "data:"
+            # which we must strip before parsing.
+            if chunk.startswith('data: '):
+              chunk = chunk[len('data: ') :]
+            yield json.loads(chunk)
+      else:
+        raise ValueError(
+            'Error parsing streaming response.'
+        )
 
   def byte_segments(self):
     if isinstance(self.byte_stream, list):
@@ -373,17 +381,22 @@ class BaseApiClient:
           if not self.project:
             self.project = project
 
-    if self._credentials.expired or not self._credentials.token:
-      # Only refresh when it needs to. Default expiration is 3600 seconds.
-      async with self._auth_lock:
-        if self._credentials.expired or not self._credentials.token:
-          # Double check that the credentials expired before refreshing.
-          await asyncio.to_thread(_refresh_auth, self._credentials)
+    if self._credentials:
+      if (
+          self._credentials.expired or not self._credentials.token
+      ):
+        # Only refresh when it needs to. Default expiration is 3600 seconds.
+        async with self._auth_lock:
+          if self._credentials.expired or not self._credentials.token:
+            # Double check that the credentials expired before refreshing.
+            await asyncio.to_thread(_refresh_auth, self._credentials)
 
-    if not self._credentials.token:
-      raise RuntimeError('Could not resolve API token from the environment')
+      if not self._credentials.token:
+        raise RuntimeError('Could not resolve API token from the environment')
 
-    return self._credentials.token
+      return self._credentials.token
+    else:
+      raise RuntimeError('Could not resolve API token from the environment')
 
   def _build_request(
       self,
@@ -428,8 +441,12 @@ class BaseApiClient:
         patched_http_options['api_version'] + '/' + path,
     )
 
-    timeout_in_seconds: Optional[Union[float, int]] = patched_http_options.get('timeout', None)
+    timeout_in_seconds: Optional[Union[float, int]] = patched_http_options.get(
+        'timeout', None
+    )
     if timeout_in_seconds:
+      # HttpOptions.timeout is in milliseconds. But httpx.Client.request()
+      # expects seconds.
       timeout_in_seconds = timeout_in_seconds / 1000.0
     else:
       timeout_in_seconds = None
@@ -503,18 +520,19 @@ class BaseApiClient:
       http_request.headers['Authorization'] = (
           f'Bearer {await self._async_access_token()}'
       )
-      if self._credentials.quota_project_id:
+      if self._credentials and self._credentials.quota_project_id:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
     if stream:
-      httpx_request = httpx.Request(
+      aclient = httpx.AsyncClient()
+      httpx_request = aclient.build_request(
           method=http_request.method,
           url=http_request.url,
           content=json.dumps(http_request.data),
           headers=http_request.headers,
+          timeout=http_request.timeout,
       )
-      aclient = httpx.AsyncClient()
       response = await aclient.send(
           httpx_request,
           stream=stream,
@@ -652,7 +670,7 @@ class BaseApiClient:
     offset = 0
     # Upload the file in chunks
     while True:
-      file_chunk = file.read(1024 * 1024 * 8)  # 8 MB chunk size
+      file_chunk = file.read(CHUNK_SIZE)
       chunk_size = 0
       if file_chunk:
         chunk_size = len(file_chunk)
@@ -679,13 +697,12 @@ class BaseApiClient:
       if upload_size <= offset:  # Status is not finalized.
         raise ValueError(
             'All content has been uploaded, but the upload status is not'
-            f' finalized. {response.headers}, body: {response.json}'
+            f' finalized.'
         )
 
     if response.headers['X-Goog-Upload-Status'] != 'final':
       raise ValueError(
-          'Failed to upload file: Upload status is not finalized. headers:'
-          f' {response.headers}, body: {response.json}'
+          'Failed to upload file: Upload status is not finalized.'
       )
     return response.json
 
@@ -708,7 +725,7 @@ class BaseApiClient:
       self,
       http_request: HttpRequest,
   ) -> HttpResponse:
-    data: str | bytes | None = None
+    data: Optional[Union[str, bytes]] = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
         data = json.dumps(http_request.data)
@@ -746,16 +763,17 @@ class BaseApiClient:
     returns:
           The response json object from the finalize request.
     """
-    return await asyncio.to_thread(
-        self.upload_file,
-        file_path,
-        upload_url,
-        upload_size,
-    )
+    if isinstance(file_path, io.IOBase):
+      return await self._async_upload_fd(file_path, upload_url, upload_size)
+    else:
+      file = anyio.Path(file_path)
+      fd = await file.open('rb')
+      async with fd:
+        return await self._async_upload_fd(fd, upload_url, upload_size)
 
   async def _async_upload_fd(
       self,
-      file: io.IOBase,
+      file: Union[io.IOBase, anyio.AsyncFile],
       upload_url: str,
       upload_size: int,
   ) -> str:
@@ -770,12 +788,45 @@ class BaseApiClient:
     returns:
           The response json object from the finalize request.
     """
-    return await asyncio.to_thread(
-        self._upload_fd,
-        file,
-        upload_url,
-        upload_size,
-    )
+    async with httpx.AsyncClient() as aclient:
+      offset = 0
+      # Upload the file in chunks
+      while True:
+        if isinstance(file, io.IOBase):
+          file_chunk = file.read(CHUNK_SIZE)
+        else:
+          file_chunk = await file.read(CHUNK_SIZE)
+        chunk_size = 0
+        if file_chunk:
+          chunk_size = len(file_chunk)
+        upload_command = 'upload'
+        # If last chunk, finalize the upload.
+        if chunk_size + offset >= upload_size:
+          upload_command += ', finalize'
+        response = await aclient.request(
+            method='POST',
+            url=upload_url,
+            content=file_chunk,
+            headers={
+                'X-Goog-Upload-Command': upload_command,
+                'X-Goog-Upload-Offset': str(offset),
+                'Content-Length': str(chunk_size),
+            },
+        )
+        offset += chunk_size
+        if response.headers.get('x-goog-upload-status') != 'active':
+          break  # upload is complete or it has been interrupted.
+
+        if upload_size <= offset:  # Status is not finalized.
+          raise ValueError(
+              'All content has been uploaded, but the upload status is not'
+              f' finalized.'
+          )
+      if response.headers.get('x-goog-upload-status') != 'final':
+        raise ValueError(
+            'Failed to upload file: Upload status is not finalized.'
+        )
+      return response.json()
 
   async def async_download_file(self, path: str, http_options):
     """Downloads the file data.
@@ -787,12 +838,31 @@ class BaseApiClient:
     returns:
           The file bytes
     """
-    return await asyncio.to_thread(
-        self.download_file,
-        path,
-        http_options,
+    http_request = self._build_request(
+        'get', path=path, request_dict={}, http_options=http_options
     )
 
+    data: Optional[Union[str, bytes]]
+    if http_request.data:
+      if not isinstance(http_request.data, bytes):
+        data = json.dumps(http_request.data)
+      else:
+        data = http_request.data
+
+    async with httpx.AsyncClient(follow_redirects=True) as aclient:
+      response = await aclient.request(
+          method=http_request.method,
+          url=http_request.url,
+          headers=http_request.headers,
+          content=data,
+          timeout=http_request.timeout,
+      )
+      errors.APIError.raise_for_response(response)
+
+      return HttpResponse(
+          response.headers, byte_stream=[response.read()]
+      ).byte_stream[0]
+
   # This method does nothing in the real api client. It is used in the
   # replay_api_client to verify the response from the SDK method matches the
   # recorded response.

google/genai/_automatic_function_calling_util.py

@@ -17,7 +17,7 @@ import inspect
 import sys
 import types as builtin_types
 import typing
-from typing import _GenericAlias, Any, Callable, get_args, get_origin, Literal, Union
+from typing import _GenericAlias, Any, Callable, get_args, get_origin, Literal, Optional, Union
 
 import pydantic
 
@@ -304,9 +304,9 @@ def _parse_schema_from_parameter(
   )
 
 
-def _get_required_fields(schema: types.Schema) -> list[str]:
+def _get_required_fields(schema: types.Schema) -> Optional[list[str]]:
   if not schema.properties:
-    return
+    return None
   return [
       field_name
       for field_name, field_schema in schema.properties.items()

google/genai/_common.py

@@ -188,6 +188,8 @@ def _remove_extra_fields(
         if isinstance(item, dict):
           _remove_extra_fields(typing.get_args(annotation)[0], item)
 
+T = typing.TypeVar('T', bound='BaseModel')
+
 
 class BaseModel(pydantic.BaseModel):
 
@@ -201,12 +203,13 @@ class BaseModel(pydantic.BaseModel):
       arbitrary_types_allowed=True,
       ser_json_bytes='base64',
       val_json_bytes='base64',
+      ignored_types=(typing.TypeVar,)
   )
 
   @classmethod
   def _from_response(
-      cls, *, response: dict[str, object], kwargs: dict[str, object]
-  ) -> 'BaseModel':
+      cls: typing.Type[T], *, response: dict[str, object], kwargs: dict[str, object]
+  ) -> T:
     # 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.

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]