google-genai 1.33.0__py3-none-any.whl → 1.53.0__py3-none-any.whl

google/genai/_api_client.py

@@ -37,13 +37,13 @@ import time
 from typing import Any, AsyncIterator, Iterator, Optional, Tuple, TYPE_CHECKING, Union
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
+import warnings
 
 import anyio
 import certifi
 import google.auth
 import google.auth.credentials
 from google.auth.credentials import Credentials
-from google.auth.transport.requests import Request
 import httpx
 from pydantic import BaseModel
 from pydantic import ValidationError
@@ -56,6 +56,7 @@ from .types import HttpOptions
 from .types import HttpOptionsOrDict
 from .types import HttpResponse as SdkHttpResponse
 from .types import HttpRetryOptions
+from .types import ResourceScope
 
 
 try:
@@ -79,7 +80,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger('google_genai._api_client')
 CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunk size
-READ_BUFFER_SIZE = 2**20
+READ_BUFFER_SIZE = 2**22
 MAX_RETRY_COUNT = 3
 INITIAL_RETRY_DELAY = 1  # second
 DELAY_MULTIPLIER = 2
@@ -196,6 +197,7 @@ def load_auth(*, project: Union[str, None]) -> Tuple[Credentials, str]:
 
 
 def refresh_auth(credentials: Credentials) -> Credentials:
+  from google.auth.transport.requests import Request
   credentials.refresh(Request())  # type: ignore[no-untyped-call]
   return credentials
 
@@ -229,7 +231,6 @@ class HttpResponse:
       headers: Union[dict[str, str], httpx.Headers, 'CIMultiDictProxy[str]'],
       response_stream: Union[Any, str] = None,
       byte_stream: Union[Any, bytes] = None,
-      session: Optional['aiohttp.ClientSession'] = None,
   ):
     if isinstance(headers, dict):
       self.headers = headers
@@ -245,7 +246,6 @@ class HttpResponse:
     self.status_code: int = 200
     self.response_stream = response_stream
     self.byte_stream = byte_stream
-    self._session = session
 
   # Async iterator for async streaming.
   def __aiter__(self) -> 'HttpResponse':
@@ -360,69 +360,76 @@ class HttpResponse:
     balance = 0
     # httpx.Response has a dedicated async line iterator.
     if isinstance(self.response_stream, httpx.Response):
-      async for line in self.response_stream.aiter_lines():
-        if not line:
-          continue
-        # In streaming mode, the response of JSON is prefixed with "data: "
-        # which we must strip before parsing.
-        if line.startswith('data: '):
-          yield line[len('data: '):]
-          continue
-
-        # When API returns an error message, it comes line by line. So we buffer
-        # the lines until a complete JSON string is read. A complete JSON string
-        # is found when the balance is 0.
-        for c in line:
-          if c == '{':
-            balance += 1
-          elif c == '}':
-            balance -= 1
-
-        chunk += line
-        if balance == 0:
+      try:
+        async for line in self.response_stream.aiter_lines():
+          if not line:
+            continue
+          # In streaming mode, the response of JSON is prefixed with "data: "
+          # which we must strip before parsing.
+          if line.startswith('data: '):
+            yield line[len('data: '):]
+            continue
+
+          # When API returns an error message, it comes line by line. So we buffer
+          # the lines until a complete JSON string is read. A complete JSON string
+          # is found when the balance is 0.
+          for c in line:
+            if c == '{':
+              balance += 1
+            elif c == '}':
+              balance -= 1
+
+          chunk += line
+          if balance == 0:
+            yield chunk
+            chunk = ''
+        # If there is any remaining chunk, yield it.
+        if chunk:
           yield chunk
-          chunk = ''
+      finally:
+        # Close the response and release the connection.
+        await self.response_stream.aclose()
 
     # aiohttp.ClientResponse uses a content stream that we read line by line.
     elif has_aiohttp and isinstance(
         self.response_stream, aiohttp.ClientResponse
     ):
-      while True:
-        # Read a line from the stream. This returns bytes.
-        line_bytes = await self.response_stream.content.readline()
-        if not line_bytes:
-          break
-        # Decode the bytes and remove trailing whitespace and newlines.
-        line = line_bytes.decode('utf-8').rstrip()
-        if not line:
-          continue
-
-        # In streaming mode, the response of JSON is prefixed with "data: "
-        # which we must strip before parsing.
-        if line.startswith('data: '):
-          yield line[len('data: '):]
-          continue
-
-        # When API returns an error message, it comes line by line. So we buffer
-        # the lines until a complete JSON string is read. A complete JSON string
-        # is found when the balance is 0.
-        for c in line:
-          if c == '{':
-            balance += 1
-          elif c == '}':
-            balance -= 1
-
-        chunk += line
-        if balance == 0:
+      try:
+        while True:
+          # Read a line from the stream. This returns bytes.
+          line_bytes = await self.response_stream.content.readline()
+          if not line_bytes:
+            break
+          # Decode the bytes and remove trailing whitespace and newlines.
+          line = line_bytes.decode('utf-8').rstrip()
+          if not line:
+            continue
+
+          # In streaming mode, the response of JSON is prefixed with "data: "
+          # which we must strip before parsing.
+          if line.startswith('data: '):
+            yield line[len('data: '):]
+            continue
+
+          # When API returns an error message, it comes line by line. So we
+          # buffer the lines until a complete JSON string is read. A complete
+          # JSON strings found when the balance is 0.
+          for c in line:
+            if c == '{':
+              balance += 1
+            elif c == '}':
+              balance -= 1
+
+          chunk += line
+          if balance == 0:
+            yield chunk
+            chunk = ''
+        # If there is any remaining chunk, yield it.
+        if chunk:
           yield chunk
-          chunk = ''
-
-    # If there is any remaining chunk, yield it.
-    if chunk:
-      yield chunk
-
-    if hasattr(self, '_session') and self._session:
-      await self._session.close()
+      finally:
+        # Release the connection back to the pool for potential reuse.
+        self.response_stream.release()
 
   @classmethod
   def _load_json_from_response(cls, response: Any) -> Any:
@@ -483,6 +490,7 @@ def retry_args(options: Optional[HttpRetryOptions]) -> _common.StringDict:
       'retry': retry,
       'reraise': True,
       'wait': wait,
+      'before_sleep': tenacity.before_sleep_log(logger, logging.INFO),
   }
 
 
@@ -540,6 +548,7 @@ class BaseApiClient:
       http_options: Optional[HttpOptionsOrDict] = None,
   ):
     self.vertexai = vertexai
+    self.custom_base_url = None
     if self.vertexai is None:
       if os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '0').lower() in [
           'true',
@@ -571,6 +580,12 @@ class BaseApiClient:
     elif isinstance(http_options, HttpOptions):
       validated_http_options = http_options
 
+    if validated_http_options.base_url_resource_scope and not validated_http_options.base_url:
+      # base_url_resource_scope is only valid when base_url is set.
+      raise ValueError(
+          'base_url must be set when base_url_resource_scope is set.'
+      )
+
     # 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)
@@ -591,7 +606,7 @@ class BaseApiClient:
     # Handle when to use Vertex AI in express mode (api key).
     # Explicit initializer arguments are already validated above.
     if self.vertexai:
-      if credentials:
+      if credentials and env_api_key:
         # Explicit credentials take precedence over implicit api_key.
         logger.info(
             'The user provided Google Cloud credentials will take precedence'
@@ -621,11 +636,20 @@ class BaseApiClient:
         )
         self.api_key = None
 
+      self.custom_base_url = (
+          validated_http_options.base_url
+          if validated_http_options.base_url
+          else None
+      )
+
+      if not self.location and not self.api_key and not self.custom_base_url:
+          self.location = 'global'
+
       # Skip fetching project from ADC if base url is provided in http options.
       if (
           not self.project
           and not self.api_key
-          and not validated_http_options.base_url
+          and not self.custom_base_url
       ):
         credentials, self.project = load_auth(project=None)
         if not self._credentials:
@@ -633,17 +657,22 @@ class BaseApiClient:
 
       has_sufficient_auth = (self.project and self.location) or self.api_key
 
-      if not has_sufficient_auth and not validated_http_options.base_url:
+      if not has_sufficient_auth and not self.custom_base_url:
         # Skip sufficient auth check if base url is provided in http options.
         raise ValueError(
-            'Project and location or API key must be set when using the Vertex '
+            'Project 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/'
-      elif validated_http_options.base_url and not has_sufficient_auth:
+      elif self.custom_base_url and not ((project and location) or api_key):
         # Avoid setting default base url and api version if base_url provided.
-        self._http_options.base_url = validated_http_options.base_url
+        # API gateway proxy can use the auth in custom headers, not url.
+        # Enable custom url if auth is not sufficient.
+        self._http_options.base_url = self.custom_base_url
+        # Clear project and location if base_url is provided.
+        self.project = None
+        self.location = None
       else:
         self._http_options.base_url = (
             f'https://{self.location}-aiplatform.googleapis.com/'
@@ -676,19 +705,79 @@ class BaseApiClient:
     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)
+    self._async_httpx_client_args = async_client_args
+
+    if self._http_options.httpx_client:
+      self._httpx_client = self._http_options.httpx_client
+    else:
+      self._httpx_client = SyncHttpxClient(**client_args)
+    if self._http_options.httpx_async_client:
+      self._async_httpx_client = self._http_options.httpx_async_client
+    else:
+      self._async_httpx_client = AsyncHttpxClient(**async_client_args)
     if self._use_aiohttp():
-      # Do it once at the genai.Client level. Share among all requests.
-      self._async_client_session_request_args = self._ensure_aiohttp_ssl_ctx(
-          self._http_options
-      )
-    self._websocket_ssl_ctx = self._ensure_websocket_ssl_ctx(self._http_options)
+      try:
+        import aiohttp  # pylint: disable=g-import-not-at-top
+        # Do it once at the genai.Client level. Share among all requests.
+        self._async_client_session_request_args = self._ensure_aiohttp_ssl_ctx(
+            self._http_options
+        )
+      except ImportError:
+        pass
+
+    # Initialize the aiohttp client session.
+    self._aiohttp_session: Optional[aiohttp.ClientSession] = None
 
     retry_kwargs = retry_args(self._http_options.retry_options)
+    self._websocket_ssl_ctx = self._ensure_websocket_ssl_ctx(self._http_options)
     self._retry = tenacity.Retrying(**retry_kwargs)
     self._async_retry = tenacity.AsyncRetrying(**retry_kwargs)
 
+  async def _get_aiohttp_session(self) -> 'aiohttp.ClientSession':
+    """Returns the aiohttp client session."""
+    if (
+        self._aiohttp_session is None
+        or self._aiohttp_session.closed
+        or self._aiohttp_session._loop.is_closed()  # pylint: disable=protected-access
+    ):
+      # Initialize the aiohttp client session if it's not set up or closed.
+      class AiohttpClientSession(aiohttp.ClientSession):  # type: ignore[misc]
+
+        def __del__(self, _warnings: Any = warnings) -> None:
+          if not self.closed:
+            context = {
+                'client_session': self,
+                'message': 'Unclosed client session',
+            }
+            if self._source_traceback is not None:
+              context['source_traceback'] = self._source_traceback
+            # Remove this self._loop.call_exception_handler(context)
+
+      class AiohttpTCPConnector(aiohttp.TCPConnector):  # type: ignore[misc]
+
+        def __del__(self, _warnings: Any = warnings) -> None:
+          if self._closed:
+            return
+          if not self._conns:
+            return
+          conns = [repr(c) for c in self._conns.values()]
+          # After v3.13.2, it may change to self._close_immediately()
+          self._close()
+          context = {
+              'connector': self,
+              'connections': conns,
+              'message': 'Unclosed connector',
+          }
+          if self._source_traceback is not None:
+            context['source_traceback'] = self._source_traceback
+          # Remove this self._loop.call_exception_handler(context)
+      self._aiohttp_session = AiohttpClientSession(
+          connector=AiohttpTCPConnector(limit=0),
+          trust_env=True,
+          read_bufsize=READ_BUFFER_SIZE,
+      )
+    return self._aiohttp_session
+
   @staticmethod
   def _ensure_httpx_ssl_ctx(
       options: HttpOptions,
@@ -767,7 +856,6 @@ class BaseApiClient:
     Returns:
       An async aiohttp ClientSession._request args.
     """
-
     verify = 'ssl'  # keep it consistent with httpx.
     async_args = options.async_client_args
     ctx = async_args.get(verify) if async_args else None
@@ -875,9 +963,15 @@ class BaseApiClient:
         has_aiohttp
         and (self._http_options.async_client_args or {}).get('transport')
         is None
+        and (self._http_options.httpx_async_client is None)
     )
 
   def _websocket_base_url(self) -> str:
+    has_sufficient_auth = (self.project and self.location) or self.api_key
+    if self.custom_base_url and not has_sufficient_auth:
+      # API gateway proxy can use the auth in custom headers, not url.
+      # Enable custom url if auth is not sufficient.
+      return self.custom_base_url
     url_parts = urlparse(self._http_options.base_url)
     return url_parts._replace(scheme='wss').geturl()  # type: ignore[arg-type, return-value]
 
@@ -992,6 +1086,11 @@ class BaseApiClient:
         and not path.startswith('projects/')
         and not query_vertex_base_models
         and (self.project or self.location)
+        and not (
+            self.custom_base_url
+            and patched_http_options.base_url_resource_scope
+            == ResourceScope.COLLECTION
+        )
     ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
 
@@ -1015,11 +1114,27 @@ class BaseApiClient:
       _common.recursive_dict_update(
           request_dict, patched_http_options.extra_body
       )
-
-    url = join_url_path(
-        base_url,
-        versioned_path,
-    )
+    url = base_url
+    if (
+        not self.custom_base_url
+        or (self.project and self.location)
+        or self.api_key
+    ):
+      if (
+          patched_http_options.base_url_resource_scope
+          == ResourceScope.COLLECTION
+      ):
+        url = join_url_path(base_url, path)
+      else:
+        url = join_url_path(
+            base_url,
+            versioned_path,
+        )
+    elif(
+        self.custom_base_url
+        and patched_http_options.base_url_resource_scope == ResourceScope.COLLECTION
+    ):
+      url = join_url_path(base_url, path)
 
     if self.api_key and self.api_key.startswith('auth_tokens/'):
       raise EphemeralTokenAPIKeyError(
@@ -1132,13 +1247,9 @@ class BaseApiClient:
 
     if stream:
       if self._use_aiohttp():
-        session = aiohttp.ClientSession(
-            headers=http_request.headers,
-            trust_env=True,
-            read_bufsize=READ_BUFFER_SIZE,
-        )
+        self._aiohttp_session = await self._get_aiohttp_session()
         try:
-          response = await session.request(
+          response = await self._aiohttp_session.request(
               method=http_request.method,
               url=http_request.url,
               headers=http_request.headers,
@@ -1159,12 +1270,8 @@ class BaseApiClient:
               self._ensure_aiohttp_ssl_ctx(self._http_options)
           )
           # Instantiate a new session with the updated SSL context.
-          session = aiohttp.ClientSession(
-              headers=http_request.headers,
-              trust_env=True,
-              read_bufsize=READ_BUFFER_SIZE,
-          )
-          response = await session.request(
+          self._aiohttp_session = await self._get_aiohttp_session()
+          response = await self._aiohttp_session.request(
               method=http_request.method,
               url=http_request.url,
               headers=http_request.headers,
@@ -1174,7 +1281,7 @@ class BaseApiClient:
           )
 
         await errors.APIError.raise_for_async_response(response)
-        return HttpResponse(response.headers, response, session=session)
+        return HttpResponse(response.headers, response)
       else:
         # aiohttp is not available. Fall back to httpx.
         httpx_request = self._async_httpx_client.build_request(
@@ -1192,22 +1299,18 @@ class BaseApiClient:
         return HttpResponse(client_response.headers, client_response)
     else:
       if self._use_aiohttp():
+        self._aiohttp_session = await self._get_aiohttp_session()
         try:
-          async with aiohttp.ClientSession(
+          response = await self._aiohttp_session.request(
+              method=http_request.method,
+              url=http_request.url,
               headers=http_request.headers,
-              trust_env=True,
-              read_bufsize=READ_BUFFER_SIZE,
-          ) 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),
-                **self._async_client_session_request_args,
-            )
-            await errors.APIError.raise_for_async_response(response)
-            return HttpResponse(response.headers, [await response.text()])
+              data=data,
+              timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+              **self._async_client_session_request_args,
+          )
+          await errors.APIError.raise_for_async_response(response)
+          return HttpResponse(response.headers, [await response.text()])
         except (
             aiohttp.ClientConnectorError,
             aiohttp.ClientConnectorDNSError,
@@ -1221,21 +1324,17 @@ class BaseApiClient:
               self._ensure_aiohttp_ssl_ctx(self._http_options)
           )
           # Instantiate a new session with the updated SSL context.
-          async with aiohttp.ClientSession(
+          self._aiohttp_session = await self._get_aiohttp_session()
+          response = await self._aiohttp_session.request(
+              method=http_request.method,
+              url=http_request.url,
               headers=http_request.headers,
-              trust_env=True,
-              read_bufsize=READ_BUFFER_SIZE,
-          ) 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),
-                **self._async_client_session_request_args,
-            )
-            await errors.APIError.raise_for_async_response(response)
-            return HttpResponse(response.headers, [await response.text()])
+              data=data,
+              timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+              **self._async_client_session_request_args,
+          )
+          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(
@@ -1305,9 +1404,21 @@ class BaseApiClient:
 
     session_response = self._request(http_request, http_options, stream=True)
     for chunk in session_response.segments():
-      yield SdkHttpResponse(
-          headers=session_response.headers, body=json.dumps(chunk)
-      )
+      chunk_dump = json.dumps(chunk)
+      try:
+        if chunk_dump.startswith('{"error":'):
+          chunk_json = json.loads(chunk_dump)
+          errors.APIError.raise_error(
+              chunk_json.get('error', {}).get('code'),
+              chunk_json,
+              session_response,
+          )
+      except json.decoder.JSONDecodeError:
+        logger.debug(
+            'Failed to decode chunk that contains an error: %s' % chunk_dump
+        )
+        pass
+      yield SdkHttpResponse(headers=session_response.headers, body=chunk_dump)
 
   async def async_request(
       self,
@@ -1341,7 +1452,21 @@ class BaseApiClient:
 
     async def async_generator():  # type: ignore[no-untyped-def]
       async for chunk in response:
-        yield SdkHttpResponse(headers=response.headers, body=json.dumps(chunk))
+        chunk_dump = json.dumps(chunk)
+        try:
+          if chunk_dump.startswith('{"error":'):
+            chunk_json = json.loads(chunk_dump)
+            await errors.APIError.raise_error_async(
+                chunk_json.get('error', {}).get('code'),
+                chunk_json,
+                response,
+            )
+        except json.decoder.JSONDecodeError:
+          logger.debug(
+              'Failed to decode chunk that contains an error: %s' % chunk_dump
+          )
+          pass
+        yield SdkHttpResponse(headers=response.headers, body=chunk_dump)
 
     return async_generator()  # type: ignore[no-untyped-call]
 
@@ -1451,7 +1576,7 @@ class BaseApiClient:
             f'All content has been uploaded, but the upload status is not'
             f' finalized.'
         )
-
+    errors.APIError.raise_for_response(response)
     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])
@@ -1551,85 +1676,81 @@ class BaseApiClient:
     offset = 0
     # Upload the file in chunks
     if self._use_aiohttp():  # pylint: disable=g-import-not-at-top
-      async with aiohttp.ClientSession(
-          headers=self._http_options.headers,
-          trust_env=True,
-          read_bufsize=READ_BUFFER_SIZE,
-      ) 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
+      self._aiohttp_session = await self._get_aiohttp_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 = (
-              http_options.get('timeout')
-              if isinstance(http_options, dict)
-              else http_options.timeout
+              self._http_options.timeout
+              if isinstance(self._http_options, dict)
+              else self._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),
-            )
+        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)
 
-            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)
+        retry_count = 0
+        response = None
+        while retry_count < MAX_RETRY_COUNT:
+          response = await self._aiohttp_session.request(
+              method='POST',
+              url=upload_url,
+              data=file_chunk,
+              headers=upload_headers,
+              timeout=aiohttp.ClientTimeout(connect=timeout_in_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 response.headers.get('X-Goog-Upload-Status'):
+            break
+          delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
+          retry_count += 1
+          await asyncio.sleep(delay_seconds)
 
-          if upload_size <= offset:  # Status is not finalized.
-            raise ValueError(
-                f'All content has been uploaded, but the upload status is not'
-                f' finalized.'
-            )
+        offset += chunk_size
         if (
             response is not None
-            and response.headers.get('X-Goog-Upload-Status') != 'final'
+            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(
-              'Failed to upload file: Upload status is not finalized.'
+              f'All content has been uploaded, but the upload status is not'
+              f' finalized.'
           )
-        return HttpResponse(
-            response.headers, response_stream=[await response.text()]
+
+      await errors.APIError.raise_for_async_response(response)
+      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()]
+      )
     else:
       # aiohttp is not available. Fall back to httpx.
       while True:
@@ -1697,6 +1818,8 @@ class BaseApiClient:
               'All content has been uploaded, but the upload status is not'
               ' finalized.'
           )
+
+      await errors.APIError.raise_for_async_response(client_response)
       if (
           client_response is not None
           and client_response.headers.get('x-goog-upload-status') != 'final'
@@ -1735,23 +1858,19 @@ class BaseApiClient:
         data = http_request.data
 
     if self._use_aiohttp():
-      async with aiohttp.ClientSession(
+      self._aiohttp_session = await self._get_aiohttp_session()
+      response = await self._aiohttp_session.request(
+          method=http_request.method,
+          url=http_request.url,
           headers=http_request.headers,
-          trust_env=True,
-          read_bufsize=READ_BUFFER_SIZE,
-      ) 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)
+          data=data,
+          timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+      )
+      await errors.APIError.raise_for_async_response(response)
 
-        return HttpResponse(
-            response.headers, byte_stream=[await 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(
@@ -1772,3 +1891,37 @@ class BaseApiClient:
   # recorded response.
   def _verify_response(self, response_model: _common.BaseModel) -> None:
     pass
+
+  def close(self) -> None:
+    """Closes the API client."""
+    # Let users close the custom client explicitly by themselves. Otherwise,
+    # close the client when the object is garbage collected.
+    if not self._http_options.httpx_client:
+      self._httpx_client.close()
+
+  async def aclose(self) -> None:
+    """Closes the API async client."""
+    # Let users close the custom client explicitly by themselves. Otherwise,
+    # close the client when the object is garbage collected.
+    if not self._http_options.httpx_async_client:
+      await self._async_httpx_client.aclose()
+    if self._aiohttp_session:
+      await self._aiohttp_session.close()
+
+  def __del__(self) -> None:
+    """Closes the API client when the object is garbage collected.
+
+    ADK uses this client so cannot rely on the genai.[Async]Client.__del__
+    for cleanup.
+    """
+
+    try:
+      if not self._http_options.httpx_client:
+        self.close()
+    except Exception:  # pylint: disable=broad-except
+      pass
+
+    try:
+      asyncio.get_running_loop().create_task(self.aclose())
+    except Exception:  # pylint: disable=broad-except
+      pass