From 9a7dbd9025daa001d551da7c1322e85f72db2f00 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 24 Oct 2025 14:32:50 +0000 Subject: [PATCH 1/4] Ensure that google-genai doesn't close httpx client provided by Pydantic AI or user --- .../pydantic_ai/providers/google.py | 29 ++++- ...est_google_httpx_client_is_not_closed.yaml | 122 ++++++++++++++++++ tests/models/test_google.py | 11 ++ 3 files changed, 161 insertions(+), 1 deletion(-) create mode 100644 tests/models/cassettes/test_google/test_google_httpx_client_is_not_closed.yaml diff --git a/pydantic_ai_slim/pydantic_ai/providers/google.py b/pydantic_ai_slim/pydantic_ai/providers/google.py index d123061d24..ffddce6ac2 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/google.py +++ b/pydantic_ai_slim/pydantic_ai/providers/google.py @@ -13,7 +13,8 @@ try: from google.auth.credentials import Credentials - from google.genai import Client + from google.genai._api_client import BaseApiClient + from google.genai.client import Client, DebugConfig from google.genai.types import HttpOptions except ImportError as _import_error: raise ImportError( @@ -186,9 +187,35 @@ def __init__( class _SafelyClosingClient(Client): + @staticmethod + def _get_api_client( + vertexai: bool | None = None, + api_key: str | None = None, + credentials: Credentials | None = None, + project: str | None = None, + location: str | None = None, + debug_config: DebugConfig | None = None, + http_options: HttpOptions | None = None, + ) -> BaseApiClient: + return _NonClosingApiClient( + vertexai=vertexai, + api_key=api_key, + credentials=credentials, + project=project, + location=location, + http_options=http_options, + ) + def close(self) -> None: # This is called from `Client.__del__`, even if `Client.__init__` raised an error before `self._api_client` is set, which would raise an `AttributeError` here. try: super().close() except AttributeError: pass + + +class _NonClosingApiClient(BaseApiClient): + async def aclose(self) -> None: + # The original implementation also calls `await self._async_httpx_client.aclose()`, but we don't want to close our `cached_async_http_client` or the one the user passed in. + if self._aiohttp_session: + await self._aiohttp_session.close() diff --git a/tests/models/cassettes/test_google/test_google_httpx_client_is_not_closed.yaml b/tests/models/cassettes/test_google/test_google_httpx_client_is_not_closed.yaml new file mode 100644 index 0000000000..9fee66b8d1 --- /dev/null +++ b/tests/models/cassettes/test_google/test_google_httpx_client_is_not_closed.yaml @@ -0,0 +1,122 @@ +interactions: +- request: + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '141' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + method: POST + parsed_body: + contents: + - parts: + - text: What is the capital of France? + role: user + generationConfig: + responseModalities: + - TEXT + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent + response: + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '549' + content-type: + - application/json; charset=UTF-8 + server-timing: + - gfet4t7; dur=581 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + parsed_body: + candidates: + - content: + parts: + - text: The capital of France is **Paris**. + role: model + finishReason: STOP + index: 0 + modelVersion: gemini-2.5-flash-lite + responseId: mI37aJyZEsGtz7IPjumZ8AM + usageMetadata: + candidatesTokenCount: 8 + promptTokenCount: 8 + promptTokensDetails: + - modality: TEXT + tokenCount: 8 + totalTokenCount: 16 + status: + code: 200 + message: OK +- request: + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '141' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + method: POST + parsed_body: + contents: + - parts: + - text: What is the capital of Mexico? + role: user + generationConfig: + responseModalities: + - TEXT + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent + response: + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '555' + content-type: + - application/json; charset=UTF-8 + server-timing: + - gfet4t7; dur=408 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + parsed_body: + candidates: + - content: + parts: + - text: The capital of Mexico is **Mexico City**. + role: model + finishReason: STOP + index: 0 + modelVersion: gemini-2.5-flash-lite + responseId: mY37aIyAFKDUz7IPv4PK4AY + usageMetadata: + candidatesTokenCount: 9 + promptTokenCount: 8 + promptTokensDetails: + - modality: TEXT + tokenCount: 8 + totalTokenCount: 17 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/test_google.py b/tests/models/test_google.py index 89150af60d..e96145e6c3 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -2923,3 +2923,14 @@ async def test_google_vertexai_image_generation(allow_model_requests: None, vert result = await agent.run('Generate an image of an axolotl.') assert result.output == snapshot(BinaryImage(data=IsBytes(), media_type='image/png', identifier='b037a4')) + + +async def test_google_httpx_client_is_not_closed(allow_model_requests: None, google_provider: GoogleProvider): + # This should not raise any errors, see https://github.com/pydantic/pydantic-ai/issues/3242. + agent = Agent('google-gla:gemini-2.5-flash-lite') + result = await agent.run('What is the capital of France?') + assert result.output == snapshot('The capital of France is **Paris**.') + + agent = Agent('google-gla:gemini-2.5-flash-lite') + result = await agent.run('What is the capital of Mexico?') + assert result.output == snapshot('The capital of Mexico is **Mexico City**.') From e634af91c38269c18dadb39f7c6f475ccd7e3b0c Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 24 Oct 2025 14:53:31 +0000 Subject: [PATCH 2/4] Link to issues --- pydantic_ai_slim/pydantic_ai/providers/google.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/google.py b/pydantic_ai_slim/pydantic_ai/providers/google.py index ffddce6ac2..fe9c6a8cc4 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/google.py +++ b/pydantic_ai_slim/pydantic_ai/providers/google.py @@ -115,7 +115,7 @@ def __init__( base_url=base_url, headers={'User-Agent': get_user_agent()}, httpx_async_client=http_client, - # TODO: Remove once https://github.com/googleapis/python-genai/pull/1509#issuecomment-3430028790 is solved. + # TODO: Remove once https://github.com/googleapis/python-genai/issues/1565 is solved. async_client_args={'transport': httpx.AsyncHTTPTransport()}, ) if not vertexai: @@ -208,6 +208,7 @@ def _get_api_client( def close(self) -> None: # This is called from `Client.__del__`, even if `Client.__init__` raised an error before `self._api_client` is set, which would raise an `AttributeError` here. + # TODO: Remove once https://github.com/googleapis/python-genai/issues/1567 is solved. try: super().close() except AttributeError: @@ -217,5 +218,6 @@ def close(self) -> None: class _NonClosingApiClient(BaseApiClient): async def aclose(self) -> None: # The original implementation also calls `await self._async_httpx_client.aclose()`, but we don't want to close our `cached_async_http_client` or the one the user passed in. + # TODO: Remove once https://github.com/googleapis/python-genai/issues/1566 is solved. if self._aiohttp_session: await self._aiohttp_session.close() From 39733a340dcef3488814f5a7ba5c3c0186614d47 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 24 Oct 2025 15:00:42 +0000 Subject: [PATCH 3/4] Fix test --- tests/models/test_google.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/models/test_google.py b/tests/models/test_google.py index e96145e6c3..4668b58c70 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -2925,12 +2925,12 @@ async def test_google_vertexai_image_generation(allow_model_requests: None, vert assert result.output == snapshot(BinaryImage(data=IsBytes(), media_type='image/png', identifier='b037a4')) -async def test_google_httpx_client_is_not_closed(allow_model_requests: None, google_provider: GoogleProvider): +async def test_google_httpx_client_is_not_closed(allow_model_requests: None, gemini_api_key: str): # This should not raise any errors, see https://github.com/pydantic/pydantic-ai/issues/3242. - agent = Agent('google-gla:gemini-2.5-flash-lite') + agent = Agent(GoogleModel('gemini-2.5-flash-lite', provider=GoogleProvider(api_key=gemini_api_key))) result = await agent.run('What is the capital of France?') assert result.output == snapshot('The capital of France is **Paris**.') - agent = Agent('google-gla:gemini-2.5-flash-lite') + agent = Agent(GoogleModel('gemini-2.5-flash-lite', provider=GoogleProvider(api_key=gemini_api_key))) result = await agent.run('What is the capital of Mexico?') assert result.output == snapshot('The capital of Mexico is **Mexico City**.') From 287bbae8599b18dbe3f7cb58e314231f767ddede Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 24 Oct 2025 15:12:31 +0000 Subject: [PATCH 4/4] coverage --- pydantic_ai_slim/pydantic_ai/providers/google.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/google.py b/pydantic_ai_slim/pydantic_ai/providers/google.py index fe9c6a8cc4..d0de4bcc4f 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/google.py +++ b/pydantic_ai_slim/pydantic_ai/providers/google.py @@ -220,4 +220,4 @@ async def aclose(self) -> None: # The original implementation also calls `await self._async_httpx_client.aclose()`, but we don't want to close our `cached_async_http_client` or the one the user passed in. # TODO: Remove once https://github.com/googleapis/python-genai/issues/1566 is solved. if self._aiohttp_session: - await self._aiohttp_session.close() + await self._aiohttp_session.close() # pragma: no cover