feat: support passing the custom aiohttp.ClientSession through HttpOptions.aiohttp_client

fixes #1662

PiperOrigin-RevId: 856374597

Author: yinghsienwu

google/genai/_api_client.py

@@ -715,19 +715,23 @@ def __init__(
       self._async_httpx_client = self._http_options.httpx_async_client
     else:
       self._async_httpx_client = AsyncHttpxClient(**async_client_args)
+
+    # Initialize the aiohttp client session.
+    self._aiohttp_session: Optional[aiohttp.ClientSession] = None
     if self._use_aiohttp():
       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
-        )
+
+        if self._http_options.aiohttp_client:
+          self._aiohttp_session = self._http_options.aiohttp_client
+        else:
+          # 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)
@@ -1892,7 +1896,7 @@ async def aclose(self) -> None:
     # 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:
+    if self._aiohttp_session and not self._http_options.aiohttp_client:
       await self._aiohttp_session.close()
 
   def __del__(self) -> None:

google/genai/tests/client/test_custom_client.py

@@ -75,3 +75,30 @@ def test_constructor_with_httpx_clients():
   assert not mldev_client.models._api_client._use_aiohttp()
 
 
+# Aiohttp
+@requires_aiohttp
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
+)
+async def test_constructor_with_aiohttp_clients():
+  api_client.has_aiohttp = True
+  mldev_http_options = {
+      'aiohttp_client': aiohttp.ClientSession(trust_env=False),
+  }
+  vertexai_http_options = {
+      'aiohttp_client': aiohttp.ClientSession(trust_env=False),
+  }
+  mldev_client = Client(
+      api_key='google_api_key', http_options=mldev_http_options
+  )
+  assert not mldev_client.models._api_client._aiohttp_session.trust_env
+
+  vertexai_client = Client(
+      vertexai=True,
+      project='fake_project_id',
+      location='fake-location',
+      http_options=vertexai_http_options,
+  )
+  assert not vertexai_client.models._api_client._aiohttp_session.trust_env
+

google/genai/tests/client/test_http_options.py

@@ -40,7 +40,7 @@ def test_patch_http_options_with_copies_all_fields():
 
   for key in http_options_keys:
     assert hasattr(patched, key)
-    if key not in ['httpx_client', 'httpx_async_client', 'aiohttp_client_session']:
+    if key not in ['httpx_client', 'httpx_async_client', 'aiohttp_client']:
       assert getattr(patched, key) is not None
   assert patched.base_url == 'https://fake-url.com/'
   assert patched.api_version == 'v1'

google/genai/types.py

@@ -108,6 +108,20 @@
     HttpxClient = None
     HttpxAsyncClient = None
 
+_is_aiohttp_imported = False
+if typing.TYPE_CHECKING:
+  from aiohttp import ClientSession
+
+  _is_aiohttp_imported = True
+else:
+  ClientSession: typing.Type = Any
+  try:
+    from aiohttp import ClientSession
+
+    _is_aiohttp_imported = True
+  except ImportError:
+    ClientSession = None
+
 logger = logging.getLogger('google_genai.types')
 _from_json_schema_warning_logged = False
 _json_schema_warning_logged = False
@@ -1936,6 +1950,10 @@ class HttpOptions(_common.BaseModel):
       default=None,
       description="""A custom httpx async client to be used for the request.""",
   )
+  aiohttp_client: Optional['ClientSession'] = Field(
+      default=None,
+      description="""A custom aiohttp client session to be used for the request.""",
+  )
 
 
 class HttpOptionsDict(TypedDict, total=False):