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

google/genai/_api_client.py

@@ -25,19 +25,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."""
@@ -85,11 +89,36 @@ def _patch_http_options(
 
 def _join_url_path(base_url: str, path: str) -> str:
   parsed_base = urlparse(base_url)
-  base_path = parsed_base.path[:-1] if parsed_base.path.endswith('/') else parsed_base.path
+  base_path = (
+      parsed_base.path[:-1]
+      if parsed_base.path.endswith('/')
+      else parsed_base.path
+  )
   path = path[1:] if path.startswith('/') else path
   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]
@@ -101,15 +130,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: 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:
@@ -124,15 +152,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
 
@@ -159,6 +187,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.
@@ -201,12 +248,14 @@ class ApiClient:
     if (project or location) and api_key:
       # API cannot consume both project/location and api_key.
       raise ValueError(
-          'Project/location and API key are mutually exclusive in the client initializer.'
+          'Project/location and API key are mutually exclusive in the client'
+          ' initializer.'
       )
     elif credentials and api_key:
       # API cannot consume both credentials and api_key.
       raise ValueError(
-          'Credentials and API key are mutually exclusive in the client initializer.'
+          'Credentials and API key are mutually exclusive in the client'
+          ' initializer.'
       )
 
     # Validate http_options if a dict is provided.
@@ -215,7 +264,7 @@ class ApiClient:
         HttpOptions.model_validate(http_options)
       except ValidationError as e:
         raise ValueError(f'Invalid http_options: {e}')
-    elif(isinstance(http_options, HttpOptions)):
+    elif isinstance(http_options, HttpOptions):
       http_options = http_options.model_dump()
 
     # Retrieve implicitly set values from the environment.
@@ -228,20 +277,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.'
         )
@@ -249,37 +302,41 @@ 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 '
             'AI API.'
         )
       if self.api_key or self.location == 'global':
-        self._http_options['base_url'] = (
-            f'https://aiplatform.googleapis.com/'
-        )
+        self._http_options['base_url'] = f'https://aiplatform.googleapis.com/'
       else:
         self._http_options['base_url'] = (
             f'https://{self.location}-aiplatform.googleapis.com/'
         )
       self._http_options['api_version'] = 'v1beta1'
-    else:  # ML Dev API
+    else:  # Implicit initialization or missing arguments.
       if not self.api_key:
-        raise ValueError('API key must be set when using the Google AI API.')
+        raise ValueError(
+            'Missing key inputs argument! To use the Google AI API,'
+            'provide (`api_key`) arguments. To use the Google Cloud API,'
+            ' provide (`vertexai`, `project` & `location`) arguments.'
+        )
       self._http_options['base_url'] = (
           'https://generativelanguage.googleapis.com/'
       )
@@ -298,6 +355,32 @@ 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,
@@ -362,8 +445,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
@@ -371,9 +456,7 @@ class ApiClient:
           http_request.method.upper(),
           http_request.url,
           headers=http_request.headers,
-          data=json.dumps(http_request.data)
-          if http_request.data
-          else None,
+          data=json.dumps(http_request.data) if http_request.data else None,
           timeout=http_request.timeout,
       )
       errors.APIError.raise_for_response(response)
@@ -413,21 +496,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,
+          data=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,
+            data=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()
@@ -449,9 +553,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(
@@ -483,8 +587,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(
@@ -503,6 +606,7 @@ class ApiClient:
     async def async_generator():
       async for chunk in response:
         yield chunk
+
     return async_generator()
 
   def upload_file(

google/genai/_api_module.py

@@ -15,6 +15,7 @@
 
 """Utilities for the API Modules of the Google Gen AI SDK."""
 
+from typing import Optional
 from . import _api_client
 
 
@@ -22,3 +23,7 @@ class BaseModule:
 
   def __init__(self, api_client_: _api_client.ApiClient):
     self._api_client = api_client_
+
+  @property
+  def vertexai(self) -> Optional[bool]:
+    return self._api_client.vertexai

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]

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,
@@ -248,7 +250,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 +270,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'

google/genai/_replay_api_client.py

@@ -60,6 +60,10 @@ def _redact_request_headers(headers):
       redacted_headers[header_name] = _redact_language_label(
           _redact_version_numbers(header_value)
       )
+    elif header_name.lower() == 'x-goog-user-project':
+      continue
+    elif header_name.lower() == 'authorization':
+      continue
     else:
       redacted_headers[header_name] = header_value
   return redacted_headers
@@ -409,6 +413,34 @@ class ReplayApiClient(ApiClient):
     else:
       return self._build_response_from_replay(http_request)
 
+  async def _async_request(
+      self,
+      http_request: HttpRequest,
+      stream: bool = False,
+  ) -> HttpResponse:
+    self._initialize_replay_session_if_not_loaded()
+    if self._should_call_api():
+      _debug_print('api mode request: %s' % http_request)
+      try:
+        result = await super()._async_request(http_request, stream)
+      except errors.APIError as e:
+        self._record_interaction(http_request, e)
+        raise e
+      if stream:
+        result_segments = []
+        async for segment in result.async_segments():
+          result_segments.append(json.dumps(segment))
+        result = HttpResponse(result.headers, result_segments)
+        self._record_interaction(http_request, result)
+        # Need to return a RecordedResponse that rebuilds the response
+        # segments since the stream has been consumed.
+      else:
+        self._record_interaction(http_request, result)
+      _debug_print('api mode result: %s' % result.json)
+      return result
+    else:
+      return self._build_response_from_replay(http_request)
+
   def upload_file(self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int):
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
@@ -453,4 +485,3 @@ class ReplayApiClient(ApiClient):
       return result
     else:
       return self._build_response_from_replay(request)
-

google/genai/_transformers.py

@@ -20,11 +20,12 @@ from collections.abc import Iterable, Mapping
 from enum import Enum, EnumMeta
 import inspect
 import io
+import logging
 import re
+import sys
 import time
 import typing
 from typing import Any, GenericAlias, Optional, Union
-import sys
 
 if typing.TYPE_CHECKING:
   import PIL.Image
@@ -34,6 +35,8 @@ import pydantic
 from . import _api_client
 from . import types
 
+logger = logging.getLogger('google_genai._transformers')
+
 if sys.version_info >= (3, 10):
   VersionedUnionType = typing.types.UnionType
   _UNION_TYPES = (typing.Union, typing.types.UnionType)
@@ -183,8 +186,15 @@ def t_extract_models(
     return response.get('tunedModels')
   elif response.get('publisherModels') is not None:
     return response.get('publisherModels')
+  elif (
+      response.get('httpHeaders') is not None
+      and response.get('jsonPayload') is None
+  ):
+    return []
   else:
-    raise ValueError('Cannot determine the models type.')
+    logger.warning('Cannot determine the models type.')
+    logger.debug('Cannot determine the models type for response: %s', response)
+    return []
 
 
 def t_caches_model(api_client: _api_client.ApiClient, model: str):
@@ -205,12 +215,17 @@ def t_caches_model(api_client: _api_client.ApiClient, model: str):
 def pil_to_blob(img) -> types.Blob:
   try:
     import PIL.PngImagePlugin
+
     PngImagePlugin = PIL.PngImagePlugin
   except ImportError:
     PngImagePlugin = None
 
   bytesio = io.BytesIO()
-  if PngImagePlugin is not None and isinstance(img, PngImagePlugin.PngImageFile) or img.mode == 'RGBA':
+  if (
+      PngImagePlugin is not None
+      and isinstance(img, PngImagePlugin.PngImageFile)
+      or img.mode == 'RGBA'
+  ):
     img.save(bytesio, format='PNG')
     mime_type = 'image/png'
   else:
@@ -249,7 +264,7 @@ def t_part(client: _api_client.ApiClient, part: PartType) -> types.Part:
 def t_parts(
     client: _api_client.ApiClient, parts: Union[list, PartType]
 ) -> list[types.Part]:
-  if parts is None:
+  if not parts:
     raise ValueError('content parts are required.')
   if isinstance(parts, list):
     return [t_part(client, part) for part in parts]
@@ -382,7 +397,10 @@ def handle_null_fields(schema: dict[str, Any]):
 def process_schema(
     schema: dict[str, Any],
     client: Optional[_api_client.ApiClient] = None,
-    defs: Optional[dict[str, Any]]=None):
+    defs: Optional[dict[str, Any]] = None,
+    *,
+    order_properties: bool = True,
+):
   """Updates the schema and each sub-schema inplace to be API-compatible.
 
   - Removes the `title` field from the schema if the client is not vertexai.
@@ -517,6 +535,16 @@ def process_schema(
         ref = defs[ref_key.split('defs/')[-1]]
         process_schema(ref, client, defs)
         properties[name] = ref
+    if (
+        len(properties.items()) > 1
+        and order_properties
+        and all(
+            ordering_key not in schema
+            for ordering_key in ['property_ordering', 'propertyOrdering']
+        )
+    ):
+      property_names = list(properties.keys())
+      schema['property_ordering'] = property_names
   elif schema_type == 'ARRAY':
     sub_schema = schema.get('items', None)
     if sub_schema is None:
@@ -539,6 +567,7 @@ def _process_enum(
           f'Enum member {member.name} value must be a string, got'
           f' {type(member.value)}'
       )
+
   class Placeholder(pydantic.BaseModel):
     placeholder: enum
 
@@ -554,7 +583,7 @@ def t_schema(
   if not origin:
     return None
   if isinstance(origin, dict):
-    process_schema(origin, client)
+    process_schema(origin, client, order_properties=False)
     return types.Schema.model_validate(origin)
   if isinstance(origin, EnumMeta):
     return _process_enum(origin, client)
@@ -563,15 +592,15 @@ def t_schema(
       # response_schema value was coerced to an empty Schema instance because it did not adhere to the Schema field annotation
       raise ValueError(f'Unsupported schema type.')
     schema = origin.model_dump(exclude_unset=True)
-    process_schema(schema, client)
+    process_schema(schema, client, order_properties=False)
     return types.Schema.model_validate(schema)
 
   if (
       # in Python 3.9 Generic alias list[int] counts as a type,
       # and breaks issubclass because it's not a class.
-      not isinstance(origin, GenericAlias) and
-      isinstance(origin, type) and
-      issubclass(origin, pydantic.BaseModel)
+      not isinstance(origin, GenericAlias)
+      and isinstance(origin, type)
+      and issubclass(origin, pydantic.BaseModel)
   ):
     schema = origin.model_json_schema()
     process_schema(schema, client)
@@ -582,6 +611,7 @@ def t_schema(
       or isinstance(origin, VersionedUnionType)
       or typing.get_origin(origin) in _UNION_TYPES
   ):
+
     class Placeholder(pydantic.BaseModel):
       placeholder: origin
 

google/genai/batches.py

@@ -15,6 +15,7 @@
 
 # Code generated by the Google Gen AI SDK generator DO NOT EDIT.
 
+import logging
 from typing import Optional, Union
 from urllib.parse import urlencode
 from . import _api_module
@@ -27,6 +28,8 @@ from ._common import get_value_by_path as getv
 from ._common import set_value_by_path as setv
 from .pagers import AsyncPager, Pager
 
+logger = logging.getLogger('google_genai.batches')
+
 
 def _BatchJobSource_to_mldev(
     api_client: ApiClient,
@@ -1050,7 +1053,7 @@ class AsyncBatches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      batch_job = client.batches.get(name='123456789')
+      batch_job = await client.aio.batches.get(name='123456789')
       print(f"Batch job: {batch_job.name}, state {batch_job.state}")
     """
 
@@ -1116,7 +1119,7 @@ class AsyncBatches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      client.batches.cancel(name='123456789')
+      await client.aio.batches.cancel(name='123456789')
     """
 
     parameter_model = types._CancelBatchJobParameters(
@@ -1222,7 +1225,7 @@ class AsyncBatches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      client.batches.delete(name='123456789')
+      await client.aio.batches.delete(name='123456789')
     """
 
     parameter_model = types._DeleteBatchJobParameters(