Skip to content

Commit f982dfb

Browse files
yinghsienwucopybara-github
yinghsienwu
authored and
copybara-github
committed
feat: Support explicitly closing the client and context manager
PiperOrigin-RevId: 809133951
1 parent dd67ade commit f982dfb

File tree

5 files changed

+339
-2
lines changed

5 files changed

+339
-2
lines changed

‎README.md‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,87 @@ from google import genai
8585
client = genai.Client()
8686
```
8787

88+
## Close a client
89+
90+
Explicitly close the sync client to ensure that resources, such as the
91+
underlying HTTP connections, are properly cleaned up and closed.
92+
93+
```python
94+
95+
from google.genai import Client
96+
97+
client = Client()
98+
response_1 = client.models.generate_content(
99+
model=MODEL_ID,
100+
contents='Hello',
101+
)
102+
response_2 = client.models.generate_content(
103+
model=MODEL_ID,
104+
contents='Ask a question',
105+
)
106+
# Close the sync client to release resources.
107+
client.close()
108+
```
109+
110+
To explicitly close the async client:
111+
112+
```python
113+
114+
from google.genai import Client
115+
116+
aclient = Client(
117+
vertexai=True, project='my-project-id', location='us-central1'
118+
).aio
119+
response_1 = await aclient.models.generate_content(
120+
model=MODEL_ID,
121+
contents='Hello',
122+
)
123+
response_2 = await aclient.models.generate_content(
124+
model=MODEL_ID,
125+
contents='Ask a question',
126+
)
127+
# Close the async client to release resources.
128+
await aclient.aclose()
129+
```
130+
131+
## Client context managers
132+
133+
By using the sync client context manager, it will close the underlying
134+
sync client when exiting the with block.
135+
136+
```python
137+
from google.genai import Client
138+
139+
with Client() as client:
140+
response_1 = client.models.generate_content(
141+
model=MODEL_ID,
142+
contents='Hello',
143+
)
144+
response_2 = client.models.generate_content(
145+
model=MODEL_ID,
146+
contents='Ask a question',
147+
)
148+
149+
```
150+
151+
By using the async client context manager, it will close the underlying
152+
async client when exiting the with block.
153+
154+
```python
155+
from google.genai import Client
156+
157+
async with Client().aio as aclient:
158+
response_1 = await aclient.models.generate_content(
159+
model=MODEL_ID,
160+
contents='Hello',
161+
)
162+
response_2 = await aclient.models.generate_content(
163+
model=MODEL_ID,
164+
contents='Ask a question',
165+
)
166+
167+
```
168+
88169
### API Selection
89170

90171
By default, the SDK uses the beta API endpoints provided by Google to support

‎google/genai/_api_client.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1763,3 +1763,14 @@ async def async_download_file(
17631763
# recorded response.
17641764
def _verify_response(self, response_model: _common.BaseModel) -> None:
17651765
pass
1766+
1767+
def close(self) -> None:
1768+
"""Closes the API client."""
1769+
self._httpx_client.close()
1770+
1771+
async def aclose(self) -> None:
1772+
"""Closes the API async client."""
1773+
1774+
await self._async_httpx_client.aclose()
1775+
if self._aiohttp_session:
1776+
await self._aiohttp_session.close()

‎google/genai/client.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#
1515

1616
import os
17+
from types import TracebackType
1718
from typing import Optional, Union
1819

1920
import google.auth
@@ -85,6 +86,44 @@ def auth_tokens(self) -> AsyncTokens:
8586
def operations(self) -> AsyncOperations:
8687
return self._operations
8788

89+
async def aclose(self) -> None:
90+
"""Closes the async client explicitly.
91+
92+
However, it doesn't close the sync client, which can be closed using the
93+
Client.close() method or using the context manager.
94+
95+
Usage:
96+
.. code-block:: python
97+
98+
from google.genai import Client
99+
100+
async_client = Client(
101+
vertexai=True, project='my-project-id', location='us-central1'
102+
).aio
103+
response_1 = await async_client.models.generate_content(
104+
model='gemini-2.0-flash',
105+
contents='Hello World',
106+
)
107+
response_2 = await async_client.models.generate_content(
108+
model='gemini-2.0-flash',
109+
contents='Hello World',
110+
)
111+
# Close the client to release resources.
112+
await async_client.aclose()
113+
"""
114+
await self._api_client.aclose()
115+
116+
async def __aenter__(self) -> 'AsyncClient':
117+
return self
118+
119+
async def __aexit__(
120+
self,
121+
exc_type: Optional[Exception],
122+
exc_value: Optional[Exception],
123+
traceback: Optional[TracebackType],
124+
) -> None:
125+
await self.aclose()
126+
88127

89128
class DebugConfig(pydantic.BaseModel):
90129
"""Configuration options that change client network behavior when testing."""
@@ -311,3 +350,41 @@ def operations(self) -> Operations:
311350
def vertexai(self) -> bool:
312351
"""Returns whether the client is using the Vertex AI API."""
313352
return self._api_client.vertexai or False
353+
354+
def close(self) -> None:
355+
"""Closes the synchronous client explicitly.
356+
357+
However, it doesn't close the async client, which can be closed using the
358+
Client.aio.aclose() method or using the async context manager.
359+
360+
Usage:
361+
.. code-block:: python
362+
363+
from google.genai import Client
364+
365+
client = Client(
366+
vertexai=True, project='my-project-id', location='us-central1'
367+
)
368+
response_1 = client.models.generate_content(
369+
model='gemini-2.0-flash',
370+
contents='Hello World',
371+
)
372+
response_2 = client.models.generate_content(
373+
model='gemini-2.0-flash',
374+
contents='Hello World',
375+
)
376+
# Close the client to release resources.
377+
client.close()
378+
"""
379+
self._api_client.close()
380+
381+
def __enter__(self) -> 'Client':
382+
return self
383+
384+
def __exit__(
385+
self,
386+
exc_type: Optional[Exception],
387+
exc_value: Optional[Exception],
388+
traceback: Optional[TracebackType],
389+
) -> None:
390+
self.close()

‎google/genai/tests/client/test_client_close.py‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
16+
"""Tests for closing the clients and context managers."""
17+
import asyncio
18+
from unittest import mock
19+
20+
from google.oauth2 import credentials
21+
import pytest
22+
try:
23+
import aiohttp
24+
AIOHTTP_NOT_INSTALLED = False
25+
except ImportError:
26+
AIOHTTP_NOT_INSTALLED = True
27+
aiohttp = mock.MagicMock()
28+
29+
30+
from ... import _api_client as api_client
31+
from ... import Client
32+
33+
34+
requires_aiohttp = pytest.mark.skipif(
35+
AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
36+
)
37+
38+
39+
def test_close_httpx_client():
40+
"""Tests that the httpx client is closed when the client is closed."""
41+
api_client.has_aiohttp = False
42+
client = Client(
43+
vertexai=True,
44+
project='test_project',
45+
location='global',
46+
)
47+
client.close()
48+
assert client._api_client._httpx_client.is_closed
49+
50+
51+
def test_httpx_client_context_manager():
52+
"""Tests that the httpx client is closed when the client is closed."""
53+
api_client.has_aiohttp = False
54+
with Client(
55+
vertexai=True,
56+
project='test_project',
57+
location='global',
58+
) as client:
59+
pass
60+
assert not client._api_client._httpx_client.is_closed
61+
62+
assert client._api_client._httpx_client.is_closed
63+
64+
65+
@pytest.mark.asyncio
66+
async def test_aclose_httpx_client():
67+
"""Tests that the httpx async client is closed when the client is closed."""
68+
api_client.has_aiohttp = False
69+
async_client = Client(
70+
vertexai=True,
71+
project='test_project',
72+
location='global',
73+
).aio
74+
await async_client.aclose()
75+
assert async_client._api_client._async_httpx_client.is_closed
76+
77+
78+
@pytest.mark.asyncio
79+
async def test_async_httpx_client_context_manager():
80+
"""Tests that the httpx async client is closed when the client is closed."""
81+
api_client.has_aiohttp = False
82+
async with Client(
83+
vertexai=True,
84+
project='test_project',
85+
location='global',
86+
).aio as async_client:
87+
pass
88+
assert not async_client._api_client._async_httpx_client.is_closed
89+
90+
assert async_client._api_client._async_httpx_client.is_closed
91+
92+
93+
@requires_aiohttp
94+
@pytest.mark.asyncio
95+
async def test_aclose_aiohttp_session():
96+
"""Tests that the aiohttp session is closed when the client is closed."""
97+
api_client.has_aiohttp = True
98+
async_client = Client(
99+
vertexai=True,
100+
project='test_project',
101+
location='global',
102+
).aio
103+
await async_client.aclose()
104+
assert async_client._api_client._aiohttp_session is None
105+
106+
107+
@requires_aiohttp
108+
@pytest.fixture
109+
def mock_request():
110+
mock_aiohttp_response = mock.Mock(spec=aiohttp.ClientSession.request)
111+
mock_aiohttp_response.return_value = mock_aiohttp_response
112+
yield mock_aiohttp_response
113+
114+
115+
def _patch_auth_default():
116+
return mock.patch(
117+
'google.auth.default',
118+
return_value=(credentials.Credentials('magic_token'), 'test_project'),
119+
autospec=True,
120+
)
121+
122+
123+
async def _aiohttp_async_response(status: int):
124+
"""Has to return a coroutine hence async."""
125+
response = mock.Mock(spec=aiohttp.ClientResponse)
126+
response.status = status
127+
response.headers = {'status-code': str(status)}
128+
response.json.return_value = {}
129+
response.text.return_value = 'test'
130+
return response
131+
132+
133+
@requires_aiohttp
134+
@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
135+
def test_aiohttp_session_context_manager(mock_request):
136+
"""Tests that the aiohttp session is closed when the client is closed."""
137+
api_client.has_aiohttp = True
138+
async def run():
139+
mock_request.side_effect = (
140+
aiohttp.ClientConnectorError(
141+
connection_key=aiohttp.client_reqrep.ConnectionKey(
142+
'localhost', 80, False, True, None, None, None
143+
),
144+
os_error=OSError,
145+
),
146+
_aiohttp_async_response(200),
147+
)
148+
with _patch_auth_default():
149+
async with Client(
150+
vertexai=True,
151+
project='test_project',
152+
location='global',
153+
).aio as async_client:
154+
# aiohttp session is created in the first request instead of client
155+
# initialization.
156+
_ = await async_client._api_client._async_request_once(
157+
api_client.HttpRequest(
158+
method='GET',
159+
url='https://example.com',
160+
headers={},
161+
data=None,
162+
timeout=None,
163+
)
164+
)
165+
assert async_client._api_client._aiohttp_session is not None
166+
assert not async_client._api_client._aiohttp_session.closed
167+
168+
assert async_client._api_client._aiohttp_session.closed
169+
170+
asyncio.run(run())

‎google/genai/tests/client/test_client_requests.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616

1717
"""Tests for client behavior when issuing requests."""
1818

19-
import http
20-
2119
from ... import _api_client as api_client
2220
from ... import Client
2321
from ... import types

0 commit comments

Comments
 (0)