google-genai 1.4.0__py3-none-any.whl → 1.6.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
@@ -28,22 +29,21 @@ import json
 import logging
 import os
 import sys
-from typing import Any, AsyncIterator, Optional, Tuple, TypedDict, Union
+from typing import Any, AsyncIterator, Optional, Tuple, Union
 from urllib.parse import urlparse, urlunparse
 import google.auth
 import google.auth.credentials
 from google.auth.credentials import Credentials
-from google.auth.transport.requests import AuthorizedSession
 from google.auth.transport.requests import Request
 import httpx
-from pydantic import BaseModel, ConfigDict, Field, ValidationError
-import requests
+from pydantic import BaseModel, Field, ValidationError
 from . import _common
 from . import errors
 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:
@@ -86,7 +86,8 @@ def _patch_http_options(
       copy_option[patch_key].update(patch_value)
     elif patch_value is not None:  # Accept empty values.
       copy_option[patch_key] = patch_value
-  _append_library_version_headers(copy_option['headers'])
+  if copy_option['headers']:
+    _append_library_version_headers(copy_option['headers'])
   return copy_option
 
 
@@ -101,7 +102,7 @@ def _join_url_path(base_url: str, path: str) -> str:
   return urlunparse(parsed_base._replace(path=base_path + '/' + path))
 
 
-def _load_auth(*, project: Union[str, None]) -> tuple[Credentials, str]:
+def _load_auth(*, project: Union[str, None]) -> Tuple[Credentials, str]:
   """Loads google auth credentials and project id."""
   credentials, loaded_project_id = google.auth.default(
       scopes=['https://www.googleapis.com/auth/cloud-platform'],
@@ -118,8 +119,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 +149,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 +202,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):
@@ -265,7 +272,9 @@ class BaseApiClient:
     validated_http_options: dict[str, Any]
     if isinstance(http_options, dict):
       try:
-        validated_http_options = HttpOptions.model_validate(http_options).model_dump()
+        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):
@@ -351,7 +360,9 @@ class BaseApiClient:
       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, validated_http_options)
+      self._http_options = _patch_http_options(
+          self._http_options, validated_http_options
+      )
     else:
       _append_library_version_headers(self._http_options['headers'])
 
@@ -359,8 +370,27 @@ class BaseApiClient:
     url_parts = urlparse(self._http_options['base_url'])
     return url_parts._replace(scheme='wss').geturl()
 
-  async def _async_access_token(self) -> str:
+  def _access_token(self) -> str:
     """Retrieves the access token for the credentials."""
+    if not self._credentials:
+      self._credentials, project = _load_auth(project=self.project)
+      if not self.project:
+        self.project = project
+
+    if self._credentials:
+      if (
+          self._credentials.expired or not self._credentials.token
+      ):
+        # Only refresh when it needs to. Default expiration is 3600 seconds.
+        _refresh_auth(self._credentials)
+      if not self._credentials.token:
+        raise RuntimeError('Could not resolve API token from the environment')
+      return self._credentials.token
+    else:
+      raise RuntimeError('Could not resolve API token from the environment')
+
+  async def _async_access_token(self) -> str:
+    """Retrieves the access token for the credentials asynchronously."""
     if not self._credentials:
       async with self._auth_lock:
         # This ensures that only one coroutine can execute the auth logic at a
@@ -373,17 +403,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,
@@ -424,12 +459,16 @@ class BaseApiClient:
     ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
     url = _join_url_path(
-        patched_http_options['base_url'],
-        patched_http_options['api_version'] + '/' + path,
+        patched_http_options.get('base_url', ''),
+        patched_http_options.get('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
@@ -447,74 +486,80 @@ class BaseApiClient:
       http_request: HttpRequest,
       stream: bool = False,
   ) -> HttpResponse:
+    data: Optional[Union[str, bytes]] = None
     if self.vertexai and not self.api_key:
-      if not self._credentials:
-        self._credentials, _ = _load_auth(project=self.project)
-      if self._credentials.quota_project_id:
+      http_request.headers['Authorization'] = (
+          f'Bearer {self._access_token()}'
+      )
+      if self._credentials and self._credentials.quota_project_id:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
-      authed_session = AuthorizedSession(self._credentials)
-      authed_session.stream = stream
-      response = authed_session.request(
-          http_request.method.upper(),
-          http_request.url,
+      data = json.dumps(http_request.data)
+    else:
+      if http_request.data:
+        if not isinstance(http_request.data, bytes):
+          data = json.dumps(http_request.data)
+        else:
+          data = http_request.data
+
+    if stream:
+      client = httpx.Client()
+      httpx_request = client.build_request(
+          method=http_request.method,
+          url=http_request.url,
+          content=data,
           headers=http_request.headers,
-          data=json.dumps(http_request.data) if http_request.data else None,
           timeout=http_request.timeout,
       )
+      response = client.send(httpx_request, stream=stream)
       errors.APIError.raise_for_response(response)
       return HttpResponse(
           response.headers, response if stream else [response.text]
       )
     else:
-      return self._request_unauthorized(http_request, stream)
-
-  def _request_unauthorized(
-      self,
-      http_request: HttpRequest,
-      stream: bool = False,
-  ) -> HttpResponse:
-    data: Optional[Union[str, bytes]] = None
-    if http_request.data:
-      if not isinstance(http_request.data, bytes):
-        data = json.dumps(http_request.data)
-      else:
-        data = http_request.data
-
-    http_session = requests.Session()
-    response = http_session.request(
-        method=http_request.method,
-        url=http_request.url,
-        headers=http_request.headers,
-        data=data,
-        timeout=http_request.timeout,
-        stream=stream,
-    )
-    errors.APIError.raise_for_response(response)
-    return HttpResponse(
-        response.headers, response if stream else [response.text]
-    )
+      with httpx.Client() as client:
+        response = client.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, response if stream else [response.text]
+        )
 
   async def _async_request(
       self, http_request: HttpRequest, stream: bool = False
   ):
-    if self.vertexai:
+    data: Optional[Union[str, bytes]] = None
+    if self.vertexai and not self.api_key:
       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
         )
+      data = json.dumps(http_request.data)
+    else:
+      if http_request.data:
+        if not isinstance(http_request.data, bytes):
+          data = json.dumps(http_request.data)
+        else:
+          data = http_request.data
+
     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),
+          content=data,
           headers=http_request.headers,
+          timeout=http_request.timeout,
       )
-      aclient = httpx.AsyncClient()
       response = await aclient.send(
           httpx_request,
           stream=stream,
@@ -529,7 +574,7 @@ class BaseApiClient:
             method=http_request.method,
             url=http_request.url,
             headers=http_request.headers,
-            content=json.dumps(http_request.data) if http_request.data else None,
+            content=data,
             timeout=http_request.timeout,
         )
         errors.APIError.raise_for_response(response)
@@ -615,7 +660,7 @@ class BaseApiClient:
 
   def upload_file(
       self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file to the given URL.
 
     Args:
@@ -637,7 +682,7 @@ class BaseApiClient:
 
   def _upload_fd(
       self, file: io.IOBase, upload_url: str, upload_size: int
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file to the given URL.
 
     Args:
@@ -652,7 +697,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)
@@ -671,7 +716,7 @@ class BaseApiClient:
           data=file_chunk,
       )
 
-      response = self._request_unauthorized(request, stream=False)
+      response = self._request(request, stream=False)
       offset += chunk_size
       if response.headers['X-Goog-Upload-Status'] != 'active':
         break  # upload is complete or it has been interrupted.
@@ -679,13 +724,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,32 +752,31 @@ 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)
       else:
         data = http_request.data
 
-    http_session = requests.Session()
-    response = http_session.request(
-        method=http_request.method,
-        url=http_request.url,
-        headers=http_request.headers,
-        data=data,
-        timeout=http_request.timeout,
-        stream=False,
-    )
+    with httpx.Client(follow_redirects=True) as client:
+      response = client.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.content])
+      errors.APIError.raise_for_response(response)
+      return HttpResponse(response.headers, byte_stream=[response.read()])
 
   async def async_upload_file(
       self,
       file_path: Union[str, io.IOBase],
       upload_url: str,
       upload_size: int,
-  ) -> str:
+  ) -> dict[str, str]:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -746,19 +789,20 @@ 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:
+  ) -> dict[str, str]:
     """Transfers a file asynchronously to the given URL.
 
     Args:
@@ -770,12 +814,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 +864,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]] = None
+    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  # type: ignore[attr-defined]
 
 import pydantic
 
@@ -41,19 +41,11 @@ _py_builtin_type_to_schema_type = {
 
 
 def _is_builtin_primitive_or_compound(
-    annotation: inspect.Parameter.annotation,
+    annotation: inspect.Parameter.annotation,  # type: ignore[valid-type]
 ) -> 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'
-        ' the Gemini API.'
-    )
-
-
 def _raise_for_default_if_mldev(schema: types.Schema):
   if schema.default is not None:
     raise ValueError(
@@ -64,12 +56,11 @@ def _raise_for_default_if_mldev(schema: types.Schema):
 
 def _raise_if_schema_unsupported(api_option: Literal['VERTEX_AI', 'GEMINI_API'], schema: types.Schema):
   if api_option == 'GEMINI_API':
-    _raise_for_any_of_if_mldev(schema)
     _raise_for_default_if_mldev(schema)
 
 
 def _is_default_value_compatible(
-    default_value: Any, annotation: inspect.Parameter.annotation
+    default_value: Any, annotation: inspect.Parameter.annotation  # type: ignore[valid-type]
 ) -> bool:
   # None type is expected to be handled external to this function
   if _is_builtin_primitive_or_compound(annotation):
@@ -292,8 +283,7 @@ def _parse_schema_from_parameter(
           ),
           func_name,
       )
-    if api_option == 'VERTEX_AI':
-      schema.required = _get_required_fields(schema)
+    schema.required = _get_required_fields(schema)
     _raise_if_schema_unsupported(api_option, schema)
     return schema
   raise ValueError(
@@ -304,9 +294,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.