google-genai 1.18.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

google/genai/_api_client.py

@@ -34,7 +34,7 @@ import ssl
 import sys
 import threading
 import time
-from typing import Any, AsyncIterator, Optional, Tuple, Union
+from typing import Any, AsyncIterator, Optional, Tuple, TYPE_CHECKING, Union
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
 
@@ -56,6 +56,19 @@ from .types import HttpOptions
 from .types import HttpOptionsDict
 from .types import HttpOptionsOrDict
 
+has_aiohttp = False
+try:
+  import aiohttp
+  has_aiohttp = True
+except ImportError:
+  pass
+
+has_aiohttp = False  # exclude from copybar export
+
+
+if TYPE_CHECKING:
+  from multidict import CIMultiDictProxy
+
 
 logger = logging.getLogger('google_genai._api_client')
 CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunk size
@@ -68,6 +81,26 @@ class EphemeralTokenAPIKeyError(ValueError):
   """Error raised when the API key is invalid."""
 
 
+# This method checks for the API key in the environment variables. Google API
+# key is precedenced over Gemini API key.
+def _get_env_api_key() -> Optional[str]:
+  """Gets the API key from environment variables, prioritizing GOOGLE_API_KEY.
+
+  Returns:
+      The API key string if found, otherwise None. Empty string is considered
+      invalid.
+  """
+  env_google_api_key = os.environ.get('GOOGLE_API_KEY', None)
+  env_gemini_api_key = os.environ.get('GEMINI_API_KEY', None)
+  if env_google_api_key and env_gemini_api_key:
+    logger.warning(
+        'Both GOOGLE_API_KEY and GEMINI_API_KEY are set. Using'
+        ' GOOGLE_API_KEY.'
+    )
+
+  return env_google_api_key or env_gemini_api_key or None
+
+
 def _append_library_version_headers(headers: dict[str, str]) -> None:
   """Appends the telemetry header to the headers dict."""
   library_label = f'google-genai-sdk/{version.__version__}'
@@ -196,7 +229,7 @@ class HttpResponse:
 
   def __init__(
       self,
-      headers: Union[dict[str, str], httpx.Headers],
+      headers: Union[dict[str, str], httpx.Headers, 'CIMultiDictProxy[str]'],
       response_stream: Union[Any, str] = None,
       byte_stream: Union[Any, bytes] = None,
   ):
@@ -262,6 +295,17 @@ class HttpResponse:
             if chunk.startswith('data: '):
               chunk = chunk[len('data: ') :]
             yield json.loads(chunk)
+      elif hasattr(self.response_stream, 'content'):
+        async for chunk in self.response_stream.content.iter_any():
+          # This is aiohttp.ClientResponse.
+          if chunk:
+            # 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)
       else:
         raise ValueError('Error parsing streaming response.')
 
@@ -371,7 +415,7 @@ class BaseApiClient:
     # Retrieve implicitly set values from the environment.
     env_project = os.environ.get('GOOGLE_CLOUD_PROJECT', None)
     env_location = os.environ.get('GOOGLE_CLOUD_LOCATION', None)
-    env_api_key = os.environ.get('GOOGLE_API_KEY', None)
+    env_api_key = _get_env_api_key()
     self.project = project or env_project
     self.location = location or env_location
     self.api_key = api_key or env_api_key
@@ -460,14 +504,17 @@ class BaseApiClient:
       if self._http_options.headers is not None:
         _append_library_version_headers(self._http_options.headers)
 
-    client_args, async_client_args = self._ensure_ssl_ctx(self._http_options)
+    client_args, async_client_args = self._ensure_httpx_ssl_ctx(
+        self._http_options
+    )
     self._httpx_client = SyncHttpxClient(**client_args)
     self._async_httpx_client = AsyncHttpxClient(**async_client_args)
 
   @staticmethod
-  def _ensure_ssl_ctx(options: HttpOptions) -> (
-      Tuple[dict[str, Any], dict[str, Any]]):
-    """Ensures the SSL context is present in the client args.
+  def _ensure_httpx_ssl_ctx(
+      options: HttpOptions,
+  ) -> Tuple[dict[str, Any], dict[str, Any]]:
+    """Ensures the SSL context is present in the HTTPX client args.
 
     Creates a default SSL context if one is not provided.
 
@@ -521,6 +568,58 @@ class BaseApiClient:
         _maybe_set(async_args, ctx),
     )
 
+  @staticmethod
+  def _ensure_aiohttp_ssl_ctx(options: HttpOptions) -> dict[str, Any]:
+    """Ensures the SSL context is present in the async client args.
+
+    Creates a default SSL context if one is not provided.
+
+    Args:
+      options: The http options to check for SSL context.
+
+    Returns:
+      An async aiohttp ClientSession._request args.
+    """
+
+    verify = 'verify'
+    async_args = options.async_client_args
+    ctx = async_args.get(verify) if async_args else None
+
+    if not ctx:
+      # Initialize the SSL context for the httpx client.
+      # Unlike requests, the aiohttp package does not automatically pull in the
+      # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be
+      # enabled explicitly. Instead of 'verify' at client level in httpx,
+      # aiohttp uses 'ssl' at request level.
+      ctx = ssl.create_default_context(
+          cafile=os.environ.get('SSL_CERT_FILE', certifi.where()),
+          capath=os.environ.get('SSL_CERT_DIR'),
+      )
+
+    def _maybe_set(
+        args: Optional[dict[str, Any]],
+        ctx: ssl.SSLContext,
+    ) -> dict[str, Any]:
+      """Sets the SSL context in the client args if not set.
+
+      Does not override the SSL context if it is already set.
+
+      Args:
+        args: The client args to to check for SSL context.
+        ctx: The SSL context to set.
+
+      Returns:
+        The client args with the SSL context included.
+      """
+      if not args or not args.get(verify):
+        args = (args or {}).copy()
+        args['ssl'] = ctx
+      else:
+        args['ssl'] = args.pop(verify)
+      return args
+
+    return _maybe_set(async_args, ctx)
+
   def _websocket_base_url(self) -> str:
     url_parts = urlparse(self._http_options.base_url)
     return url_parts._replace(scheme='wss').geturl()  # type: ignore[arg-type, return-value]
@@ -717,33 +816,96 @@ class BaseApiClient:
           data = http_request.data
 
     if stream:
-      httpx_request = self._async_httpx_client.build_request(
-          method=http_request.method,
-          url=http_request.url,
-          content=data,
-          headers=http_request.headers,
-          timeout=http_request.timeout,
-      )
-      response = await self._async_httpx_client.send(
-          httpx_request,
-          stream=stream,
-      )
-      await errors.APIError.raise_for_async_response(response)
-      return HttpResponse(
-          response.headers, response if stream else [response.text]
-      )
+      if has_aiohttp:
+        session = aiohttp.ClientSession(
+            headers=http_request.headers,
+        )
+        if self._http_options.async_client_args:
+          # When using aiohttp request options with ssl context, the latency will higher than using httpx.
+          # Use it only if necessary. Otherwise, httpx asyncclient is faster.
+          async_client_args = self._ensure_aiohttp_ssl_ctx(
+              self._http_options
+          )
+          response = await session.request(
+              method=http_request.method,
+              url=http_request.url,
+              headers=http_request.headers,
+              data=data,
+              timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+              **async_client_args,
+          )
+        else:
+          # Aiohttp performs better than httpx w/o ssl context.
+          response = await session.request(
+              method=http_request.method,
+              url=http_request.url,
+              headers=http_request.headers,
+              data=data,
+              timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+          )
+        await errors.APIError.raise_for_async_response(response)
+        return HttpResponse(response.headers, response)
+      else:
+        # aiohttp is not available. Fall back to httpx.
+        httpx_request = self._async_httpx_client.build_request(
+            method=http_request.method,
+            url=http_request.url,
+            content=data,
+            headers=http_request.headers,
+            timeout=http_request.timeout,
+        )
+        client_response = await self._async_httpx_client.send(
+            httpx_request,
+            stream=stream,
+        )
+        await errors.APIError.raise_for_async_response(client_response)
+        return HttpResponse(client_response.headers, client_response)
     else:
-      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,
-      )
-      await errors.APIError.raise_for_async_response(response)
-      return HttpResponse(
-          response.headers, response if stream else [response.text]
-      )
+      if has_aiohttp:
+        if self._http_options.async_client_args:
+          # Note that when using aiohttp request options with ssl context, the
+          # latency will higher than using httpx async client with ssl context.
+          async_client_args = self._ensure_aiohttp_ssl_ctx(
+              self._http_options
+          )
+          async with aiohttp.ClientSession(
+              headers=http_request.headers
+          ) as session:
+            response = await session.request(
+                method=http_request.method,
+                url=http_request.url,
+                headers=http_request.headers,
+                data=data,
+                timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+                **async_client_args,
+            )
+            await errors.APIError.raise_for_async_response(response)
+            return HttpResponse(response.headers, [await response.text()])
+        else:
+          # Aiohttp performs better than httpx if not using ssl context.
+          async with aiohttp.ClientSession(
+              headers=http_request.headers
+          ) as session:
+            response = await session.request(
+                method=http_request.method,
+                url=http_request.url,
+                headers=http_request.headers,
+                data=data,
+                timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+            )
+            await errors.APIError.raise_for_async_response(response)
+            return HttpResponse(response.headers, [await response.text()])
+      else:
+        # aiohttp is not available. Fall back to httpx.
+        client_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,
+        )
+        await errors.APIError.raise_for_async_response(client_response)
+        return HttpResponse(client_response.headers, [client_response.text])
 
   def get_read_only_http_options(self) -> dict[str, Any]:
     if isinstance(self._http_options, BaseModel):
@@ -1028,68 +1190,155 @@ class BaseApiClient:
     """
     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'
-      http_options = http_options if http_options else self._http_options
-      timeout = (
-          http_options.get('timeout')
-          if isinstance(http_options, dict)
-          else http_options.timeout
-      )
-      if timeout is None:
-        # Per request timeout is not configured. Check the global timeout.
-        timeout = (
-            self._http_options.timeout
-            if isinstance(self._http_options, dict)
-            else self._http_options.timeout
+    if has_aiohttp:  # pylint: disable=g-import-not-at-top
+      async with aiohttp.ClientSession(
+          headers=self._http_options.headers
+      ) as session:
+        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'
+          http_options = http_options if http_options else self._http_options
+          timeout = (
+              http_options.get('timeout')
+              if isinstance(http_options, dict)
+              else http_options.timeout
+          )
+          if timeout is None:
+            # Per request timeout is not configured. Check the global timeout.
+            timeout = (
+                self._http_options.timeout
+                if isinstance(self._http_options, dict)
+                else self._http_options.timeout
+            )
+          timeout_in_seconds = _get_timeout_in_seconds(timeout)
+          upload_headers = {
+              'X-Goog-Upload-Command': upload_command,
+              'X-Goog-Upload-Offset': str(offset),
+              'Content-Length': str(chunk_size),
+          }
+          _populate_server_timeout_header(upload_headers, timeout_in_seconds)
+
+          retry_count = 0
+          response = None
+          while retry_count < MAX_RETRY_COUNT:
+            response = await session.request(
+                method='POST',
+                url=upload_url,
+                data=file_chunk,
+                headers=upload_headers,
+                timeout=aiohttp.ClientTimeout(connect=timeout_in_seconds),
+            )
+
+            if response.headers.get('X-Goog-Upload-Status'):
+              break
+            delay_seconds = INITIAL_RETRY_DELAY * (
+                DELAY_MULTIPLIER**retry_count
+            )
+            retry_count += 1
+            time.sleep(delay_seconds)
+
+          offset += chunk_size
+          if (
+              response is not None
+              and 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(
+                f'All content has been uploaded, but the upload status is not'
+                f' finalized.'
+            )
+        if (
+            response is not None
+            and response.headers.get('X-Goog-Upload-Status') != 'final'
+        ):
+          raise ValueError(
+              'Failed to upload file: Upload status is not finalized.'
+          )
+        return HttpResponse(
+            response.headers, response_stream=[await response.text()]
         )
-      timeout_in_seconds = _get_timeout_in_seconds(timeout)
-      upload_headers = {
-          'X-Goog-Upload-Command': upload_command,
-          'X-Goog-Upload-Offset': str(offset),
-          'Content-Length': str(chunk_size),
-      }
-      _populate_server_timeout_header(upload_headers, timeout_in_seconds)
-
-      retry_count = 0
-      while retry_count < MAX_RETRY_COUNT:
-        response = await self._async_httpx_client.request(
-            method='POST',
-            url=upload_url,
-            content=file_chunk,
-            headers=upload_headers,
-            timeout=timeout_in_seconds,
+    else:
+      # aiohttp is not available. Fall back to httpx.
+      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'
+        http_options = http_options if http_options else self._http_options
+        timeout = (
+            http_options.get('timeout')
+            if isinstance(http_options, dict)
+            else http_options.timeout
         )
-        if response.headers.get('x-goog-upload-status'):
-          break
-        delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
-        retry_count += 1
-        time.sleep(delay_seconds)
-
-      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.
+        if timeout is None:
+          # Per request timeout is not configured. Check the global timeout.
+          timeout = (
+              self._http_options.timeout
+              if isinstance(self._http_options, dict)
+              else self._http_options.timeout
+          )
+        timeout_in_seconds = _get_timeout_in_seconds(timeout)
+        upload_headers = {
+            'X-Goog-Upload-Command': upload_command,
+            'X-Goog-Upload-Offset': str(offset),
+            'Content-Length': str(chunk_size),
+        }
+        _populate_server_timeout_header(upload_headers, timeout_in_seconds)
+
+        retry_count = 0
+        client_response = None
+        while retry_count < MAX_RETRY_COUNT:
+          client_response = await self._async_httpx_client.request(
+              method='POST',
+              url=upload_url,
+              content=file_chunk,
+              headers=upload_headers,
+              timeout=timeout_in_seconds,
+          )
+          if client_response is not None and client_response.headers and client_response.headers.get('x-goog-upload-status'):
+            break
+          delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
+          retry_count += 1
+          time.sleep(delay_seconds)
+
+        offset += chunk_size
+        if (
+            client_response is not None
+            and client_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'
+              ' finalized.'
+          )
+      if (
+          client_response is not None
+          and client_response.headers.get('x-goog-upload-status') != 'final'
+      ):
         raise ValueError(
-            'All content has been uploaded, but the upload status is not'
-            f' finalized.'
+            'Failed to upload file: Upload status is not finalized.'
         )
-    if response.headers.get('x-goog-upload-status') != 'final':
-      raise ValueError(
-          'Failed to upload file: Upload status is not finalized.'
-      )
-    return HttpResponse(response.headers, response_stream=[response.text])
+      return HttpResponse(client_response.headers, response_stream=[client_response.text])
 
   async def async_download_file(
       self,
@@ -1117,18 +1366,34 @@ class BaseApiClient:
       else:
         data = http_request.data
 
-    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,
-    )
-    await errors.APIError.raise_for_async_response(response)
+    if has_aiohttp:
+      async with aiohttp.ClientSession(headers=http_request.headers) as session:
+        response = await session.request(
+            method=http_request.method,
+            url=http_request.url,
+            headers=http_request.headers,
+            data=data,
+            timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+        )
+        await errors.APIError.raise_for_async_response(response)
 
-    return HttpResponse(
-        response.headers, byte_stream=[response.read()]
-    ).byte_stream[0]
+        return HttpResponse(
+            response.headers, byte_stream=[await response.read()]
+        ).byte_stream[0]
+    else:
+      # aiohttp is not available. Fall back to httpx.
+      client_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,
+      )
+      await errors.APIError.raise_for_async_response(client_response)
+
+      return HttpResponse(
+          client_response.headers, byte_stream=[client_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/_common.py

@@ -253,7 +253,21 @@ class BaseModel(pydantic.BaseModel):
     # To maintain forward compatibility, we need to remove extra fields from
     # the response.
     # We will provide another mechanism to allow users to access these fields.
-    _remove_extra_fields(cls, response)
+
+    # For Agent Engine we don't want to call _remove_all_fields because the
+    # user may pass a dict that is not a subclass of BaseModel.
+    # If more modules require we skip this, we may want a different approach
+    should_skip_removing_fields = (
+        kwargs is not None and
+        'config' in kwargs and
+        kwargs['config'] is not None and
+        isinstance(kwargs['config'], dict) and
+        'include_all_fields' in kwargs['config']
+        and kwargs['config']['include_all_fields']
+    )
+
+    if not should_skip_removing_fields:
+      _remove_extra_fields(cls, response)
     validated_response = cls.model_validate(response)
     return validated_response
 

google/genai/_live_converters.py

@@ -981,7 +981,13 @@ def _Tool_to_vertex(
     )
 
   if getv(from_object, ['url_context']) is not None:
-    raise ValueError('url_context parameter is not supported in Vertex AI.')
+    setv(
+        to_object,
+        ['urlContext'],
+        _UrlContext_to_vertex(
+            api_client, getv(from_object, ['url_context']), to_object
+        ),
+    )
 
   if getv(from_object, ['code_execution']) is not None:
     setv(to_object, ['codeExecution'], getv(from_object, ['code_execution']))

google/genai/_replay_api_client.py

@@ -454,12 +454,7 @@ class ReplayApiClient(BaseApiClient):
     if isinstance(response_model, list):
       response_model = response_model[0]
     print('response_model: ', response_model.model_dump(exclude_none=True))
-    if isinstance(response_model, GenerateVideosOperation):
-      actual = response_model.model_dump(
-          exclude={'result'}, exclude_none=True, mode='json'
-      )
-    else:
-      actual = response_model.model_dump(exclude_none=True, mode='json')
+    actual = response_model.model_dump(exclude_none=True, mode='json')
     expected = interaction.response.sdk_response_segments[
         self._sdk_response_index
     ]

google/genai/_tokens_converters.py

@@ -981,7 +981,13 @@ def _Tool_to_vertex(
     )
 
   if getv(from_object, ['url_context']) is not None:
-    raise ValueError('url_context parameter is not supported in Vertex AI.')
+    setv(
+        to_object,
+        ['urlContext'],
+        _UrlContext_to_vertex(
+            api_client, getv(from_object, ['url_context']), to_object
+        ),
+    )
 
   if getv(from_object, ['code_execution']) is not None:
     setv(to_object, ['codeExecution'], getv(from_object, ['code_execution']))

google/genai/batches.py

@@ -18,6 +18,7 @@
 import logging
 from typing import Any, Optional, Union
 from urllib.parse import urlencode
+
 from . import _api_module
 from . import _common
 from . import _extra_utils

google/genai/caches.py

@@ -18,6 +18,7 @@
 import logging
 from typing import Any, Optional, Union
 from urllib.parse import urlencode
+
 from . import _api_module
 from . import _common
 from . import _transformers as t
@@ -1092,7 +1093,13 @@ def _Tool_to_vertex(
     )
 
   if getv(from_object, ['url_context']) is not None:
-    raise ValueError('url_context parameter is not supported in Vertex AI.')
+    setv(
+        to_object,
+        ['urlContext'],
+        _UrlContext_to_vertex(
+            api_client, getv(from_object, ['url_context']), to_object
+        ),
+    )
 
   if getv(from_object, ['code_execution']) is not None:
     setv(to_object, ['codeExecution'], getv(from_object, ['code_execution']))