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

google/genai/_api_client.py

@@ -14,8 +14,12 @@
 #
 
 
-"""Base client for calling HTTP APIs sending and receiving JSON."""
+"""Base client for calling HTTP APIs sending and receiving JSON.
 
+The BaseApiClient is intended to be a private module and is subject to change.
+"""
+
+import anyio
 import asyncio
 import copy
 from dataclasses import dataclass
@@ -41,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:
@@ -65,7 +70,7 @@ def _append_library_version_headers(headers: dict[str, str]) -> None:
 
 
 def _patch_http_options(
-    options: HttpOptionsDict, patch_options: HttpOptionsDict
+    options: HttpOptionsDict, patch_options: dict[str, Any]
 ) -> HttpOptionsDict:
   # use shallow copy so we don't override the original objects.
   copy_option = HttpOptionsDict()
@@ -115,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
@@ -131,7 +137,7 @@ class HttpRequest:
 # TODO(b/394358912): Update this class to use a SDKResponse class that can be
 # generated and used for all languages.
 class BaseResponse(_common.BaseModel):
-  http_headers: dict[str, str] = Field(
+  http_headers: Optional[dict[str, str]] = Field(
       default=None, description='The http headers of the response.'
   )
 
@@ -144,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,
   ):
@@ -197,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):
@@ -224,17 +235,17 @@ class HttpResponse:
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
 
 
-class ApiClient:
+class BaseApiClient:
   """Client for calling HTTP APIs sending and receiving JSON."""
 
   def __init__(
       self,
-      vertexai: Union[bool, None] = None,
-      api_key: Union[str, None] = None,
-      credentials: google.auth.credentials.Credentials = None,
-      project: Union[str, None] = None,
-      location: Union[str, None] = None,
-      http_options: HttpOptionsOrDict = None,
+      vertexai: Optional[bool] = None,
+      api_key: Optional[str] = None,
+      credentials: Optional[google.auth.credentials.Credentials] = None,
+      project: Optional[str] = None,
+      location: Optional[str] = None,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ):
     self.vertexai = vertexai
     if self.vertexai is None:
@@ -258,14 +269,15 @@ class ApiClient:
           ' initializer.'
       )
 
-    # Validate http_options if a dict is provided.
+    # Validate http_options if it is provided.
+    validated_http_options: dict[str, Any]
     if isinstance(http_options, dict):
       try:
-        HttpOptions.model_validate(http_options)
+        validated_http_options = HttpOptions.model_validate(http_options).model_dump()
       except ValidationError as e:
         raise ValueError(f'Invalid http_options: {e}')
     elif isinstance(http_options, HttpOptions):
-      http_options = http_options.model_dump()
+      validated_http_options = http_options.model_dump()
 
     # Retrieve implicitly set values from the environment.
     env_project = os.environ.get('GOOGLE_CLOUD_PROJECT', None)
@@ -347,7 +359,7 @@ class ApiClient:
       self._http_options['headers']['x-goog-api-key'] = self.api_key
     # Update the http options with the user provided http options.
     if http_options:
-      self._http_options = _patch_http_options(self._http_options, http_options)
+      self._http_options = _patch_http_options(self._http_options, validated_http_options)
     else:
       _append_library_version_headers(self._http_options['headers'])
 
@@ -369,24 +381,29 @@ class ApiClient:
           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,
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsOrDict = None,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> HttpRequest:
     # Remove all special dict keys such as _url and _query.
     keys_to_delete = [key for key in request_dict.keys() if key.startswith('_')]
@@ -424,8 +441,12 @@ class ApiClient:
         patched_http_options['api_version'] + '/' + path,
     )
 
-    timeout_in_seconds = 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
@@ -471,7 +492,7 @@ class ApiClient:
       http_request: HttpRequest,
       stream: bool = False,
   ) -> HttpResponse:
-    data = None
+    data: Optional[Union[str, bytes]] = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
         data = json.dumps(http_request.data)
@@ -499,18 +520,19 @@ class ApiClient:
       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,
-          data=json.dumps(http_request.data),
+          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,
@@ -525,7 +547,7 @@ class ApiClient:
             method=http_request.method,
             url=http_request.url,
             headers=http_request.headers,
-            data=json.dumps(http_request.data) if http_request.data else None,
+            content=json.dumps(http_request.data) if http_request.data else None,
             timeout=http_request.timeout,
         )
         errors.APIError.raise_for_response(response)
@@ -545,7 +567,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsOrDict = None,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -563,7 +585,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: Optional[HttpOptionsDict] = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -578,7 +600,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: Optional[HttpOptionsOrDict] = None,
   ) -> dict[str, object]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -595,7 +617,7 @@ class ApiClient:
       http_method: str,
       path: str,
       request_dict: dict[str, object],
-      http_options: HttpOptionsDict = None,
+      http_options: Optional[HttpOptionsDict] = None,
   ):
     http_request = self._build_request(
         http_method, path, request_dict, http_options
@@ -648,7 +670,7 @@ class ApiClient:
     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)
@@ -675,13 +697,12 @@ class ApiClient:
       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
 
@@ -704,10 +725,10 @@ class ApiClient:
       self,
       http_request: HttpRequest,
   ) -> HttpResponse:
-    data = None
+    data: Optional[Union[str, bytes]] = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
-        data = json.dumps(http_request.data, cls=RequestJsonEncoder)
+        data = json.dumps(http_request.data)
       else:
         data = http_request.data
 
@@ -742,16 +763,17 @@ class ApiClient:
     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:
@@ -766,12 +788,45 @@ class ApiClient:
     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.
@@ -783,14 +838,33 @@ class ApiClient:
     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.
-  def _verify_response(self, response_model: BaseModel):
+  def _verify_response(self, response_model: _common.BaseModel):
     pass

google/genai/_api_module.py

@@ -21,7 +21,7 @@ from . import _api_client
 
 class BaseModule:
 
-  def __init__(self, api_client_: _api_client.ApiClient):
+  def __init__(self, api_client_: _api_client.BaseApiClient):
     self._api_client = api_client_
 
   @property

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
 
@@ -31,12 +31,12 @@ else:
   VersionedUnionType = typing._UnionGenericAlias
 
 _py_builtin_type_to_schema_type = {
-    str: 'STRING',
-    int: 'INTEGER',
-    float: 'NUMBER',
-    bool: 'BOOLEAN',
-    list: 'ARRAY',
-    dict: 'OBJECT',
+    str: types.Type.STRING,
+    int: types.Type.INTEGER,
+    float: types.Type.NUMBER,
+    bool: types.Type.BOOLEAN,
+    list: types.Type.ARRAY,
+    dict: types.Type.OBJECT,
 }
 
 
@@ -145,7 +145,7 @@ def _parse_schema_from_parameter(
           for arg in get_args(param.annotation)
       )
   ):
-    schema.type = 'OBJECT'
+    schema.type = _py_builtin_type_to_schema_type[dict]
     schema.any_of = []
     unique_types = set()
     for arg in get_args(param.annotation):
@@ -183,7 +183,7 @@ def _parse_schema_from_parameter(
     origin = get_origin(param.annotation)
     args = get_args(param.annotation)
     if origin is dict:
-      schema.type = 'OBJECT'
+      schema.type = _py_builtin_type_to_schema_type[dict]
       if param.default is not inspect.Parameter.empty:
         if not _is_default_value_compatible(param.default, param.annotation):
           raise ValueError(default_value_error_msg)
@@ -195,7 +195,7 @@ def _parse_schema_from_parameter(
         raise ValueError(
             f'Literal type {param.annotation} must be a list of strings.'
         )
-      schema.type = 'STRING'
+      schema.type = _py_builtin_type_to_schema_type[str]
       schema.enum = list(args)
       if param.default is not inspect.Parameter.empty:
         if not _is_default_value_compatible(param.default, param.annotation):
@@ -204,7 +204,7 @@ def _parse_schema_from_parameter(
       _raise_if_schema_unsupported(api_option, schema)
       return schema
     if origin is list:
-      schema.type = 'ARRAY'
+      schema.type = _py_builtin_type_to_schema_type[list]
       schema.items = _parse_schema_from_parameter(
           api_option,
           inspect.Parameter(
@@ -222,7 +222,7 @@ def _parse_schema_from_parameter(
       return schema
     if origin is Union:
       schema.any_of = []
-      schema.type = 'OBJECT'
+      schema.type = _py_builtin_type_to_schema_type[dict]
       unique_types = set()
       for arg in args:
         # The first check is for NoneType in Python 3.9, since the __name__
@@ -280,7 +280,7 @@ def _parse_schema_from_parameter(
         and param.default is not None
     ):
       schema.default = param.default
-    schema.type = 'OBJECT'
+    schema.type = _py_builtin_type_to_schema_type[dict]
     schema.properties = {}
     for field_name, field_info in param.annotation.model_fields.items():
       schema.properties[field_name] = _parse_schema_from_parameter(
@@ -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.
@@ -266,7 +269,7 @@ def encode_unserializable_types(data: dict[str, object]) -> dict[str, object]:
     A dictionary with json.dumps() incompatible type (e.g. bytes datetime)
     to compatible type (e.g. base64 encoded string, isoformat date string).
   """
-  processed_data = {}
+  processed_data: dict[str, object] = {}
   if not isinstance(data, dict):
     return data
   for key, value in data.items():