google-genai 1.2.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."""
@@ -94,6 +98,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 +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:
@@ -128,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
 
@@ -163,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.
@@ -234,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.'
         )
@@ -255,20 +302,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 '
@@ -306,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,
@@ -370,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
@@ -419,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()
@@ -455,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(
@@ -489,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(

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,6 +20,7 @@ from collections.abc import Iterable, Mapping
 from enum import Enum, EnumMeta
 import inspect
 import io
+import logging
 import re
 import sys
 import time
@@ -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):
@@ -254,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]

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(

google/genai/caches.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
@@ -26,6 +27,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.caches')
+
 
 def _Part_to_mldev(
     api_client: ApiClient,
@@ -1183,7 +1186,7 @@ class Caches(_api_module.BaseModule):
     .. code-block:: python
 
       contents = ... // Initialize the content to cache.
-      response = await client.aio.caches.create(
+      response = client.caches.create(
           model= ... // The publisher model id
           contents=contents,
           config={
@@ -1251,8 +1254,7 @@ class Caches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      await client.aio.caches.get(name= ... ) // The server-generated resource
-      name.
+      client.caches.get(name= ... ) // The server-generated resource name.
     """
 
     parameter_model = types._GetCachedContentParameters(
@@ -1314,8 +1316,7 @@ class Caches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      await client.aio.caches.delete(name= ... ) // The server-generated
-      resource name.
+      client.caches.delete(name= ... ) // The server-generated resource name.
     """
 
     parameter_model = types._DeleteCachedContentParameters(
@@ -1377,7 +1378,7 @@ class Caches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      response = await client.aio.caches.update(
+      response = client.caches.update(
           name= ... // The server-generated resource name.
           config={
               'ttl': '7600s',
@@ -1439,8 +1440,8 @@ class Caches(_api_module.BaseModule):
 
     .. code-block:: python
 
-      cached_contents = await client.aio.caches.list(config={'page_size': 2})
-      async for cached_content in cached_contents:
+      cached_contents = client.caches.list(config={'page_size': 2})
+      for cached_content in cached_contents:
         print(cached_content)
     """
 

google/genai/client.py

@@ -27,6 +27,7 @@ from .chats import AsyncChats, Chats
 from .files import AsyncFiles, Files
 from .live import AsyncLive
 from .models import AsyncModels, Models
+from .operations import AsyncOperations, Operations
 from .tunings import AsyncTunings, Tunings
 
 
@@ -42,6 +43,7 @@ class AsyncClient:
     self._batches = AsyncBatches(self._api_client)
     self._files = AsyncFiles(self._api_client)
     self._live = AsyncLive(self._api_client)
+    self._operations = AsyncOperations(self._api_client)
 
   @property
   def models(self) -> AsyncModels:
@@ -71,6 +73,9 @@ class AsyncClient:
   def live(self) -> AsyncLive:
     return self._live
 
+  @property
+  def operations(self) -> AsyncOperations:
+    return self._operations
 
 class DebugConfig(pydantic.BaseModel):
   """Configuration options that change client network behavior when testing."""
@@ -100,9 +105,9 @@ class Client:
   `api_key="your-api-key"` or by defining `GOOGLE_API_KEY="your-api-key"` as an
   environment variable
 
-  Vertex AI API users can provide inputs argument as `vertexai=false,
+  Vertex AI API users can provide inputs argument as `vertexai=True,
   project="your-project-id", location="us-central1"` or by defining
-  `GOOGLE_GENAI_USE_VERTEXAI=false`, `GOOGLE_CLOUD_PROJECT` and
+  `GOOGLE_GENAI_USE_VERTEXAI=true`, `GOOGLE_CLOUD_PROJECT` and
   `GOOGLE_CLOUD_LOCATION` environment variables.
 
   Attributes:
@@ -205,6 +210,7 @@ class Client:
     self._caches = Caches(self._api_client)
     self._batches = Batches(self._api_client)
     self._files = Files(self._api_client)
+    self._operations = Operations(self._api_client)
 
   @staticmethod
   def _get_api_client(
@@ -270,7 +276,11 @@ class Client:
   def files(self) -> Files:
     return self._files
 
+  @property
+  def operations(self) -> Operations:
+    return self._operations
+
   @property
   def vertexai(self) -> bool:
     """Returns whether the client is using the Vertex AI API."""
-    return self._api_client.vertexai or False
+    return self._api_client.vertexai or False

google/genai/errors.py

@@ -16,7 +16,8 @@
 """Error classes for the GenAI SDK."""
 
 from typing import Any, Optional, TYPE_CHECKING, Union
-
+import httpx
+import json
 import requests
 
 
@@ -34,7 +35,9 @@ class APIError(Exception):
   response: Optional[Any] = None
 
   def __init__(
-      self, code: int, response: Union[requests.Response, 'ReplayResponse']
+      self,
+      code: int,
+      response: Union[requests.Response, 'ReplayResponse', httpx.Response],
   ):
     self.response = response
 
@@ -48,6 +51,18 @@ class APIError(Exception):
             'message': response.text,
             'status': response.reason,
         }
+    elif isinstance(response, httpx.Response):
+      try:
+        response_json = response.json()
+      except (json.decoder.JSONDecodeError, httpx.ResponseNotRead):
+        try:
+          message = response.text
+        except httpx.ResponseNotRead:
+          message = None
+        response_json = {
+            'message': message,
+            'status': response.reason_phrase,
+        }
     else:
       response_json = response.body_segments[0].get('error', {})
 
@@ -89,7 +104,7 @@ class APIError(Exception):
 
   @classmethod
   def raise_for_response(
-      cls, response: Union[requests.Response, 'ReplayResponse']
+      cls, response: Union[requests.Response, 'ReplayResponse', httpx.Response]
   ):
     """Raises an error with detailed error message if the response has an error status."""
     if response.status_code == 200: