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

google/genai/_api_client.py

@@ -29,16 +29,14 @@ 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
@@ -88,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
 
 
@@ -103,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'],
@@ -189,9 +188,11 @@ class HttpResponse:
         if chunk:
           # In streaming mode, the chunk of JSON is prefixed with "data:" which
           # we must strip before parsing.
-          if chunk.startswith(b'data: '):
-            chunk = chunk[len(b'data: ') :]
-          yield json.loads(str(chunk, 'utf-8'))
+          if not isinstance(chunk, str):
+            chunk = chunk.decode('utf-8')
+          if chunk.startswith('data: '):
+            chunk = chunk[len('data: ') :]
+          yield json.loads(chunk)
 
   async def async_segments(self) -> AsyncIterator[Any]:
     if isinstance(self.response_stream, list):
@@ -207,8 +208,10 @@ class HttpResponse:
         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.
+            # In async streaming mode, the chunk of JSON is prefixed with
+            # "data:" which we must strip before parsing.
+            if not isinstance(chunk, str):
+              chunk = chunk.decode('utf-8')
             if chunk.startswith('data: '):
               chunk = chunk[len('data: ') :]
             yield json.loads(chunk)
@@ -235,6 +238,41 @@ class HttpResponse:
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
 
 
+class SyncHttpxClient(httpx.Client):
+  """Sync httpx client."""
+
+  def __init__(self, **kwargs: Any) -> None:
+    """Initializes the httpx client."""
+    kwargs.setdefault('follow_redirects', True)
+    super().__init__(**kwargs)
+
+  def __del__(self) -> None:
+    """Closes the httpx client."""
+    if self.is_closed:
+      return
+    try:
+      self.close()
+    except Exception:
+      pass
+
+
+class AsyncHttpxClient(httpx.AsyncClient):
+  """Async httpx client."""
+
+  def __init__(self, **kwargs: Any) -> None:
+    """Initializes the httpx client."""
+    kwargs.setdefault('follow_redirects', True)
+    super().__init__(**kwargs)
+
+  def __del__(self) -> None:
+    if self.is_closed:
+      return
+    try:
+      asyncio.get_running_loop().create_task(self.aclose())
+    except Exception:
+      pass
+
+
 class BaseApiClient:
   """Client for calling HTTP APIs sending and receiving JSON."""
 
@@ -273,7 +311,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):
@@ -359,16 +399,40 @@ 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'])
+    # Initialize the httpx client.
+    self._httpx_client = SyncHttpxClient()
+    self._async_httpx_client = AsyncHttpxClient()
 
   def _websocket_base_url(self):
     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
@@ -437,8 +501,8 @@ 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(
@@ -464,59 +528,54 @@ 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) if http_request.data else None
+    else:
+      if http_request.data:
+        if not isinstance(http_request.data, bytes):
+          data = json.dumps(http_request.data) if http_request.data else None
+        else:
+          data = http_request.data
+
+    if stream:
+      httpx_request = self._httpx_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 = self._httpx_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]
-    )
+      response = self._httpx_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()}'
       )
@@ -524,16 +583,23 @@ class BaseApiClient:
         http_request.headers['x-goog-user-project'] = (
             self._credentials.quota_project_id
         )
+      data = json.dumps(http_request.data) if http_request.data else None
+    else:
+      if http_request.data:
+        if not isinstance(http_request.data, bytes):
+          data = json.dumps(http_request.data) if http_request.data else None
+        else:
+          data = http_request.data
+
     if stream:
-      aclient = httpx.AsyncClient()
-      httpx_request = aclient.build_request(
+      httpx_request = self._async_httpx_client.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,
       )
-      response = await aclient.send(
+      response = await self._async_httpx_client.send(
           httpx_request,
           stream=stream,
       )
@@ -542,18 +608,17 @@ class BaseApiClient:
           response.headers, response if stream else [response.text]
       )
     else:
-      async with httpx.AsyncClient() as aclient:
-        response = await aclient.request(
-            method=http_request.method,
-            url=http_request.url,
-            headers=http_request.headers,
-            content=json.dumps(http_request.data) if http_request.data else None,
-            timeout=http_request.timeout,
-        )
-        errors.APIError.raise_for_response(response)
-        return HttpResponse(
-            response.headers, response if stream else [response.text]
-        )
+      response = await self._async_httpx_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]
+      )
 
   def get_read_only_http_options(self) -> HttpOptionsDict:
     copied = HttpOptionsDict()
@@ -633,7 +698,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:
@@ -655,7 +720,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:
@@ -678,7 +743,7 @@ class BaseApiClient:
       # If last chunk, finalize the upload.
       if chunk_size + offset >= upload_size:
         upload_command += ', finalize'
-      request = HttpRequest(
+      response = self._httpx_client.request(
           method='POST',
           url=upload_url,
           headers={
@@ -686,25 +751,22 @@ class BaseApiClient:
               'X-Goog-Upload-Offset': str(offset),
               'Content-Length': str(chunk_size),
           },
-          data=file_chunk,
+          content=file_chunk,
       )
-
-      response = self._request_unauthorized(request, stream=False)
       offset += chunk_size
-      if response.headers['X-Goog-Upload-Status'] != 'active':
+      if response.headers['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['X-Goog-Upload-Status'] != 'final':
+    if response.headers['x-goog-upload-status'] != 'final':
       raise ValueError(
           'Failed to upload file: Upload status is not finalized.'
       )
-    return response.json
+    return response.json()
 
   def download_file(self, path: str, http_options):
     """Downloads the file data.
@@ -719,12 +781,7 @@ class BaseApiClient:
     http_request = self._build_request(
         'get', path=path, request_dict={}, http_options=http_options
     )
-    return self._download_file_request(http_request).byte_stream[0]
 
-  def _download_file_request(
-      self,
-      http_request: HttpRequest,
-  ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
     if http_request.data:
       if not isinstance(http_request.data, bytes):
@@ -732,25 +789,25 @@ class BaseApiClient:
       else:
         data = http_request.data
 
-    http_session = requests.Session()
-    response = http_session.request(
+    response = self._httpx_client.request(
         method=http_request.method,
         url=http_request.url,
         headers=http_request.headers,
-        data=data,
+        content=data,
         timeout=http_request.timeout,
-        stream=False,
     )
 
     errors.APIError.raise_for_response(response)
-    return HttpResponse(response.headers, byte_stream=[response.content])
+    return HttpResponse(
+        response.headers, byte_stream=[response.read()]
+    ).byte_stream[0]
 
   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:
@@ -776,7 +833,7 @@ class BaseApiClient:
       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:
@@ -788,45 +845,44 @@ class BaseApiClient:
     returns:
           The response json object from the finalize request.
     """
-    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':
+    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 self._async_httpx_client.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(
-            'Failed to upload file: Upload status is not finalized.'
+            'All content has been uploaded, but the upload status is not'
+            f' finalized.'
         )
-      return response.json()
+    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.
@@ -842,26 +898,25 @@ class BaseApiClient:
         'get', path=path, request_dict={}, http_options=http_options
     )
 
-    data: Optional[Union[str, bytes]]
+    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)
+    response = await self._async_httpx_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.read()]
-      ).byte_stream[0]
+    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

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, Optional, 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(

google/genai/_common.py

@@ -20,7 +20,7 @@ import datetime
 import enum
 import functools
 import typing
-from typing import Union
+from typing import Any, Union
 import uuid
 import warnings
 
@@ -93,7 +93,7 @@ def set_value_by_path(data, keys, value):
     data[keys[-1]] = value
 
 
-def get_value_by_path(data: object, keys: list[str]):
+def get_value_by_path(data: Any, keys: list[str]):
   """Examples:
 
   get_value_by_path({'a': {'b': v}}, ['a', 'b'])
@@ -128,7 +128,7 @@ def get_value_by_path(data: object, keys: list[str]):
   return data
 
 
-def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
+def convert_to_dict(obj: object) -> Any:
   """Recursively converts a given object to a dictionary.
 
   If the object is a Pydantic model, it uses the model's `model_dump()` method.
@@ -137,7 +137,9 @@ def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
     obj: The object to convert.
 
   Returns:
-    A dictionary representation of the object.
+    A dictionary representation of the object, a list of objects if a list is
+    passed, or the object itself if it is not a dictionary, list, or Pydantic
+    model.
   """
   if isinstance(obj, pydantic.BaseModel):
     return obj.model_dump(exclude_none=True)
@@ -150,7 +152,7 @@ def convert_to_dict(obj: dict[str, object]) -> dict[str, object]:
 
 
 def _remove_extra_fields(
-    model: pydantic.BaseModel, response: dict[str, object]
+    model: Any, response: dict[str, object]
 ) -> None:
   """Removes extra fields from the response that are not in the model.
 

google/genai/_replay_api_client.py

@@ -520,11 +520,14 @@ class ReplayApiClient(BaseApiClient):
     else:
       return self._build_response_from_replay(request).json
 
-  def _download_file_request(self, request):
+  def download_file(self, path: str, http_options: HttpOptions):
     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 = super()._download_file_request(request)
+        result = super().download_file(path, http_options)
       except HTTPError as e:
         result = HttpResponse(
             e.response.headers, [json.dumps({'reason': e.response.reason})]
@@ -534,7 +537,7 @@ class ReplayApiClient(BaseApiClient):
       self._record_interaction(request, result)
       return result
     else:
-      return self._build_response_from_replay(request)
+      return self._build_response_from_replay(request).byte_stream[0]
 
   async def async_download_file(self, path: str, http_options):
     self._initialize_replay_session_if_not_loaded()