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

google/genai/_api_client.py

@@ -14,7 +14,10 @@
 #
 
 
-"""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 asyncio
 import copy
@@ -25,19 +28,23 @@ import json
 import logging
 import os
 import sys
-from typing import Any, Optional, Tuple, TypedDict, Union
+from typing import Any, AsyncIterator, Optional, Tuple, TypedDict, 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 . import _common
 from . import errors
 from . import version
 from .types import HttpOptions, HttpOptionsDict, HttpOptionsOrDict
 
+logger = logging.getLogger('google_genai._api_client')
+
 
 def _append_library_version_headers(headers: dict[str, str]) -> None:
   """Appends the telemetry header to the headers dict."""
@@ -61,7 +68,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()
@@ -94,6 +101,27 @@ 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]:
+  """Loads google auth credentials and project id."""
+  credentials, loaded_project_id = google.auth.default(
+      scopes=['https://www.googleapis.com/auth/cloud-platform'],
+  )
+
+  if not project:
+    project = loaded_project_id
+
+  if not project:
+    raise ValueError(
+        'Could not resolve project using application default credentials.'
+    )
+
+  return credentials, project
+
+
+def _refresh_auth(credentials: Credentials) -> None:
+  credentials.refresh(Request())
+
+
 @dataclass
 class HttpRequest:
   headers: dict[str, str]
@@ -105,15 +133,14 @@ class HttpRequest:
 
 # TODO(b/394358912): Update this class to use a SDKResponse class that can be
 # generated and used for all languages.
-@dataclass
-class BaseResponse:
-  http_headers: dict[str, str]
+class BaseResponse(_common.BaseModel):
+  http_headers: Optional[dict[str, str]] = Field(
+      default=None, description='The http headers of the response.'
+  )
 
-  @property
-  def dict(self) -> dict[str, Any]:
-    if isinstance(self, dict):
-      return self
-    return {'httpHeaders': self.http_headers}
+  json_payload: Optional[Any] = Field(
+      default=None, description='The json payload of the response.'
+  )
 
 
 class HttpResponse:
@@ -128,15 +155,15 @@ class HttpResponse:
     self.headers = headers
     self.response_stream = response_stream
     self.byte_stream = byte_stream
-    self.segment_iterator = self.segments()
 
   # Async iterator for async streaming.
   def __aiter__(self):
+    self.segment_iterator = self.async_segments()
     return self
 
   async def __anext__(self):
     try:
-      return next(self.segment_iterator)
+      return await self.segment_iterator.__anext__()
     except StopIteration:
       raise StopAsyncIteration
 
@@ -163,6 +190,25 @@ class HttpResponse:
             chunk = chunk[len(b'data: ') :]
           yield json.loads(str(chunk, 'utf-8'))
 
+  async def async_segments(self) -> AsyncIterator[Any]:
+    if isinstance(self.response_stream, list):
+      # list of objects retrieved from replay or from non-streaming API.
+      for chunk in self.response_stream:
+        yield json.loads(chunk) if chunk else {}
+    elif self.response_stream is None:
+      async for c in []:
+        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)
+
   def byte_segments(self):
     if isinstance(self.byte_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
@@ -181,17 +227,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:
@@ -215,14 +261,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)
@@ -234,20 +281,24 @@ class ApiClient:
 
     self._credentials = credentials
     self._http_options = HttpOptionsDict()
+    # Initialize the lock. This lock will be used to protect access to the
+    # credentials. This is crucial for thread safety when multiple coroutines
+    # might be accessing the credentials at the same time.
+    self._auth_lock = asyncio.Lock()
 
     # Handle when to use Vertex AI in express mode (api key).
     # Explicit initializer arguments are already validated above.
     if self.vertexai:
       if credentials:
         # Explicit credentials take precedence over implicit api_key.
-        logging.info(
+        logger.info(
             'The user provided Google Cloud credentials will take precedence'
             + ' over the API key from the environment variable.'
         )
         self.api_key = None
       elif (env_location or env_project) and api_key:
         # Explicit api_key takes precedence over implicit project/location.
-        logging.info(
+        logger.info(
             'The user provided Vertex AI API key will take precedence over the'
             + ' project/location from the environment variables.'
         )
@@ -255,20 +306,22 @@ class ApiClient:
         self.location = None
       elif (project or location) and env_api_key:
         # Explicit project/location takes precedence over implicit api_key.
-        logging.info(
+        logger.info(
             'The user provided project/location will take precedence over the'
             + ' Vertex AI API key from the environment variable.'
         )
         self.api_key = None
       elif (env_location or env_project) and env_api_key:
         # Implicit project/location takes precedence over implicit api_key.
-        logging.info(
+        logger.info(
             'The project/location from the environment variables will take'
             + ' precedence over the API key from the environment variables.'
         )
         self.api_key = None
       if not self.project and not self.api_key:
-        self.project = google.auth.default()[1]
+        credentials, self.project = _load_auth(project=None)
+        if not self._credentials:
+          self._credentials = credentials
       if not ((self.project and self.location) or self.api_key):
         raise ValueError(
             'Project and location or API key must be set when using the Vertex '
@@ -298,7 +351,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'])
 
@@ -306,12 +359,38 @@ class ApiClient:
     url_parts = urlparse(self._http_options['base_url'])
     return url_parts._replace(scheme='wss').geturl()
 
+  async def _async_access_token(self) -> str:
+    """Retrieves the access token for the credentials."""
+    if not self._credentials:
+      async with self._auth_lock:
+        # This ensures that only one coroutine can execute the auth logic at a
+        # time for thread safety.
+        if not self._credentials:
+          # Double check that the credentials are not set before loading them.
+          self._credentials, project = await asyncio.to_thread(
+              _load_auth, project=self.project
+          )
+          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 not self._credentials.token:
+      raise RuntimeError('Could not resolve API token from the environment')
+
+    return self._credentials.token
+
   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('_')]
@@ -349,7 +428,7 @@ 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:
       timeout_in_seconds = timeout_in_seconds / 1000.0
     else:
@@ -370,8 +449,10 @@ class ApiClient:
   ) -> HttpResponse:
     if self.vertexai and not self.api_key:
       if not self._credentials:
-        self._credentials, _ = google.auth.default(
-            scopes=['https://www.googleapis.com/auth/cloud-platform'],
+        self._credentials, _ = _load_auth(project=self.project)
+      if 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
@@ -394,7 +475,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)
@@ -419,21 +500,42 @@ class ApiClient:
       self, http_request: HttpRequest, stream: bool = False
   ):
     if self.vertexai:
-      if not self._credentials:
-        self._credentials, _ = google.auth.default(
-            scopes=['https://www.googleapis.com/auth/cloud-platform'],
+      http_request.headers['Authorization'] = (
+          f'Bearer {await self._async_access_token()}'
+      )
+      if self._credentials.quota_project_id:
+        http_request.headers['x-goog-user-project'] = (
+            self._credentials.quota_project_id
         )
-      return await asyncio.to_thread(
-          self._request,
-          http_request,
-          stream=stream,
+    if stream:
+      httpx_request = httpx.Request(
+          method=http_request.method,
+          url=http_request.url,
+          content=json.dumps(http_request.data),
+          headers=http_request.headers,
       )
-    else:
-      return await asyncio.to_thread(
-          self._request,
-          http_request,
+      aclient = httpx.AsyncClient()
+      response = await aclient.send(
+          httpx_request,
           stream=stream,
       )
+      errors.APIError.raise_for_response(response)
+      return HttpResponse(
+          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]
+        )
 
   def get_read_only_http_options(self) -> HttpOptionsDict:
     copied = HttpOptionsDict()
@@ -447,7 +549,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
@@ -455,9 +557,9 @@ class ApiClient:
     response = self._request(http_request, stream=False)
     json_response = response.json
     if not json_response:
-      base_response = BaseResponse(response.headers).dict
-      return base_response
-
+      return BaseResponse(http_headers=response.headers).model_dump(
+          by_alias=True
+      )
     return json_response
 
   def request_streamed(
@@ -465,7 +567,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
@@ -480,7 +582,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
@@ -489,8 +591,7 @@ class ApiClient:
     result = await self._async_request(http_request=http_request, stream=False)
     json_response = result.json
     if not json_response:
-      base_response = BaseResponse(result.headers).dict
-      return base_response
+      return BaseResponse(http_headers=result.headers).model_dump(by_alias=True)
     return json_response
 
   async def async_request_streamed(
@@ -498,7 +599,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
@@ -607,10 +708,10 @@ class ApiClient:
       self,
       http_request: HttpRequest,
   ) -> HttpResponse:
-    data = None
+    data: str | bytes | None = 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
 
@@ -695,5 +796,5 @@ class ApiClient:
   # 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

@@ -15,10 +15,15 @@
 
 """Utilities for the API Modules of the Google Gen AI SDK."""
 
+from typing import Optional
 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
+  def vertexai(self) -> Optional[bool]:
+    return self._api_client.vertexai

google/genai/_automatic_function_calling_util.py

@@ -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(

google/genai/_common.py

@@ -60,6 +60,12 @@ def set_value_by_path(data, keys, value):
         for d in data[key_name]:
           set_value_by_path(d, keys[i + 1 :], value)
       return
+    elif key.endswith('[0]'):
+      key_name = key[:-3]
+      if key_name not in data:
+        data[key_name] = [{}]
+      set_value_by_path(data[key_name][0], keys[i + 1 :], value)
+      return
 
     data = data.setdefault(key, {})
 
@@ -106,6 +112,12 @@ def get_value_by_path(data: object, keys: list[str]):
         return [get_value_by_path(d, keys[i + 1 :]) for d in data[key_name]]
       else:
         return None
+    elif key.endswith('[0]'):
+      key_name = key[:-3]
+      if key_name in data and data[key_name]:
+        return get_value_by_path(data[key_name][0], keys[i + 1 :])
+      else:
+        return None
     else:
       if key in data:
         data = data[key]
@@ -193,7 +205,7 @@ class BaseModel(pydantic.BaseModel):
 
   @classmethod
   def _from_response(
-      cls, response: dict[str, object], kwargs: dict[str, object]
+      cls, *, response: dict[str, object], kwargs: dict[str, object]
   ) -> 'BaseModel':
     # To maintain forward compatibility, we need to remove extra fields from
     # the response.
@@ -254,7 +266,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():

google/genai/_extra_utils.py

@@ -34,6 +34,8 @@ else:
 
 _DEFAULT_MAX_REMOTE_CALLS_AFC = 10
 
+logger = logging.getLogger('google_genai.models')
+
 
 def format_destination(
     src: str,
@@ -74,7 +76,7 @@ def get_function_map(
       if config and isinstance(config, dict)
       else config
   )
-  function_map = {}
+  function_map: dict[str, object] = {}
   if not config_model:
     return function_map
   if config_model.tools:
@@ -218,15 +220,16 @@ def get_function_response_parts(
     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:
-      response = {'result': invoke_function_from_dict_args(args, func)}
+      func_response = {'result': invoke_function_from_dict_args(args, func)}
     except Exception as e:  # pylint: disable=broad-except
-      response = {'error': str(e)}
-    func_response = types.Part.from_function_response(
-        name=func_name, response=response
+      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)
+    func_response_parts.append(func_response_part)
   return func_response_parts
 
 
@@ -248,7 +251,7 @@ def should_disable_afc(
       is not None
       and int(config_model.automatic_function_calling.maximum_remote_calls) <= 0
   ):
-    logging.warning(
+    logger.warning(
         'max_remote_calls in automatic_function_calling_config'
         f' {config_model.automatic_function_calling.maximum_remote_calls} is'
         ' less than or equal to 0. Disabling automatic function calling.'
@@ -268,9 +271,12 @@ def should_disable_afc(
       config_model.automatic_function_calling.disable
       and config_model.automatic_function_calling.maximum_remote_calls
       is not None
+      # exclude the case where max_remote_calls is set to 10 by default.
+      and 'maximum_remote_calls'
+      in config_model.automatic_function_calling.model_fields_set
       and int(config_model.automatic_function_calling.maximum_remote_calls) > 0
   ):
-    logging.warning(
+    logger.warning(
         '`automatic_function_calling.disable` is set to `True`. And'
         ' `automatic_function_calling.maximum_remote_calls` is a'
         ' positive number'