From 5f2fa237ef4f31487ae920ada33587ec253b2a43 Mon Sep 17 00:00:00 2001 From: James McMinn Date: Tue, 16 Jun 2026 23:31:04 +0100 Subject: [PATCH 1/2] fix: reuse Gemini client instead of building one per call Each run_prompt/run_prompt_async call constructed a fresh genai.Client and never closed it. The client owns httpx connection pools, so under sustained load the abandoned pools accumulated and process memory grew steadily. - Cache the sync client per (project, location) and reuse it; the sync client is thread-safe and the set of pairs is tiny, so this is a bounded process-lifetime singleton. - run_prompt now calls the sync API directly, dropping the per-call asyncio.run event loop as well. - run_prompt_async still builds a client per call (its async transport binds to the running loop, which is often short-lived), but now closes it in a finally so the pool is released. - Extract _build_request / _parse_response so the sync and async paths share request building and response handling. Co-Authored-By: Claude Opus 4.8 --- src/genai_utils/gemini.py | 209 ++++++++++++++++++++----------- tests/genai_utils/test_gemini.py | 71 ++++++++++- 2 files changed, 207 insertions(+), 73 deletions(-) diff --git a/src/genai_utils/gemini.py b/src/genai_utils/gemini.py index d27141f..552dadb 100644 --- a/src/genai_utils/gemini.py +++ b/src/genai_utils/gemini.py @@ -1,7 +1,7 @@ -import asyncio import logging import os import re +from functools import lru_cache from typing import Any import requests @@ -343,6 +343,101 @@ def get_thinking_config( return None +@lru_cache(maxsize=None) +def _get_client(project: str, location: str) -> genai.Client: + """Return a process-wide Vertex AI client for a project/location pair. + + The client owns HTTP connection pools, so it is built once and reused for + every call rather than constructed per request. A fresh client each call is + never closed, so its pools accumulate and memory grows steadily under load. + google-genai's sync client is safe to share across threads, and the set of + distinct (project, location) pairs is tiny, so caching every pair for the + process lifetime is bounded.""" + return genai.Client(vertexai=True, project=project, location=location) + + +def _build_request( + prompt: str, + video_uri: str | None, + output_schema: types.SchemaUnion | None, + system_instruction: str | None, + generation_config: dict[str, Any], + safety_settings: list[types.SafetySetting], + model_config: ModelConfig, + use_grounding: bool, + do_thinking: bool, + inline_citations: bool, + labels: dict[str, str], +) -> dict[str, Any]: + """Build the keyword arguments for a `models.generate_content` call. + + Shared by the sync and async entry points so they stay in lockstep.""" + # make a copy of the generation config so it doesn't change between runs + built_gen_config = {**generation_config} + + # construct the input, adding the video if provided + parts = [] + if video_uri: + parts.append(types.Part.from_uri(file_uri=video_uri, mime_type="video/mp4")) + parts.append(types.Part.from_text(text=prompt)) + + # define the schema for the output of the model + if output_schema: + built_gen_config["response_mime_type"] = "application/json" + built_gen_config["response_schema"] = output_schema + + # sort out grounding if required + if use_grounding: + if output_schema: + raise GeminiError( + "You cannot use structured output and grounding together." + ) + grounding_tool = types.Tool(google_search=types.GoogleSearch()) + built_gen_config["tools"] = [grounding_tool] + + if inline_citations and not use_grounding: + raise GeminiError("Inline citations only work if `use_grounding = True`") + merged_labels = validate_labels(DEFAULT_LABELS | labels) + + return { + "model": model_config.model_name, + "contents": types.Content(role="user", parts=parts), + "config": types.GenerateContentConfig( + system_instruction=system_instruction, + safety_settings=safety_settings, + **built_gen_config, + labels=merged_labels, + thinking_config=get_thinking_config(model_config.model_name, do_thinking), + ), + } + + +def _parse_response( + response: types.GenerateContentResponse, + use_grounding: bool, + inline_citations: bool, +) -> str: + """Validate a generate_content response and return its text.""" + if not (response.candidates and response.text and isinstance(response.text, str)): + raise GeminiError( + f"No model output: possible reason: {response.prompt_feedback}" + ) + + if use_grounding: + grounding_ran = check_grounding_ran(response) + if not grounding_ran: + _logger.error( + "Grounding Info: GROUNDING FAILED - see previous log messages for reason" + ) + raise NoGroundingError("Grounding did not run") + + if inline_citations and response.candidates[0].grounding_metadata: + text_with_citations = add_citations(response) + return text_with_citations + + return response.text + + def run_prompt( prompt: str, video_uri: str | None = None, @@ -424,21 +519,25 @@ class Movie(BaseModel): .. _safety settings: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters .. _grounding: https://ai.google.dev/gemini-api/docs/google-search """ - return asyncio.run( - run_prompt_async( - prompt=prompt, - video_uri=video_uri, - output_schema=output_schema, - system_instruction=system_instruction, - generation_config=generation_config, - safety_settings=safety_settings, - model_config=model_config, - use_grounding=use_grounding, - do_thinking=do_thinking, - inline_citations=inline_citations, - labels=labels, - ) + if model_config is None: + model_config = generate_model_config() + + request = _build_request( + prompt, + video_uri, + output_schema, + system_instruction, + generation_config, + safety_settings, + model_config, + use_grounding, + do_thinking, + inline_citations, + labels, ) + client = _get_client(model_config.project, model_config.location) + response = client.models.generate_content(**request) + return _parse_response(response, use_grounding, inline_citations) async def run_prompt_async( @@ -522,69 +621,35 @@ class Movie(BaseModel): .. _safety settings: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters .. _grounding: https://ai.google.dev/gemini-api/docs/google-search """ - # make a copy of the generation config so it doesn't change between runs - built_gen_config = {**generation_config} if model_config is None: model_config = generate_model_config() + request = _build_request( + prompt, + video_uri, + output_schema, + system_instruction, + generation_config, + safety_settings, + model_config, + use_grounding, + do_thinking, + inline_citations, + labels, + ) + + # The async transport binds to the running event loop and this is often + # invoked under a short-lived loop (asyncio.run), so unlike the sync path we + # can't safely share one client across calls. Build one here and close its + # async transport so the connection pool isn't leaked. client = genai.Client( vertexai=True, project=model_config.project, location=model_config.location, ) + try: + response = await client.aio.models.generate_content(**request) + finally: + await client.aio.aclose() - # construct the input, adding the video if provided - parts = [] - if video_uri: - parts.append(types.Part.from_uri(file_uri=video_uri, mime_type="video/mp4")) - - parts.append(types.Part.from_text(text=prompt)) - - # define the schema for the output of the model - if output_schema: - built_gen_config["response_mime_type"] = "application/json" - built_gen_config["response_schema"] = output_schema - - # sort out grounding if required - if use_grounding: - if output_schema: - raise GeminiError( - "You cannot use structured output and grounding together." - ) - grounding_tool = types.Tool(google_search=types.GoogleSearch()) - built_gen_config["tools"] = [grounding_tool] - - if inline_citations and not use_grounding: - raise GeminiError("Inline citations only work if `use_grounding = True`") - merged_labels = validate_labels(DEFAULT_LABELS | labels) - - response = await client.aio.models.generate_content( - model=model_config.model_name, - contents=types.Content(role="user", parts=parts), - config=types.GenerateContentConfig( - system_instruction=system_instruction, - safety_settings=safety_settings, - **built_gen_config, - labels=merged_labels, - thinking_config=get_thinking_config(model_config.model_name, do_thinking), - ), - ) - - if not (response.candidates and response.text and isinstance(response.text, str)): - raise GeminiError( - f"No model output: possible reason: {response.prompt_feedback}" - ) - - if use_grounding: - grounding_ran = check_grounding_ran(response) - if not grounding_ran: - _logger.error( - "Grounding Info: GROUNDING FAILED - see previous log messages for reason" - ) - raise NoGroundingError("Grounding did not run") - - if inline_citations and response.candidates[0].grounding_metadata: - text_with_citations = add_citations(response) - return text_with_citations - - return response.text + return _parse_response(response, use_grounding, inline_citations) diff --git a/tests/genai_utils/test_gemini.py b/tests/genai_utils/test_gemini.py index 1df6e07..aff083b 100644 --- a/tests/genai_utils/test_gemini.py +++ b/tests/genai_utils/test_gemini.py @@ -1,5 +1,5 @@ import os -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, patch from google.genai import Client, types from google.genai.client import AsyncClient @@ -12,8 +12,10 @@ GeminiError, ModelConfig, NoGroundingError, + _get_client, generate_model_config, get_thinking_config, + run_prompt, run_prompt_async, validate_labels, ) @@ -67,6 +69,7 @@ async def test_dont_overwrite_generation_config(mock_client): client = Mock(Client) models = Mock(Models) async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() models.generate_content.return_value = get_dummy() client.aio = async_client @@ -100,6 +103,7 @@ async def test_error_if_grounding_with_schema(mock_client): client = Mock(Client) models = Mock(Models) async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() models.generate_content.return_value = get_dummy() client.aio = async_client @@ -127,6 +131,7 @@ async def test_error_if_citations_and_no_grounding(mock_client): client = Mock(Client) models = Mock(Models) async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() models.generate_content.return_value = get_dummy() client.aio = async_client @@ -154,6 +159,7 @@ async def test_no_grounding_error_when_grounding_does_not_run(mock_client): client = Mock(Client) models = Mock(Models) async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() async def get_no_grounding_metadata_response(): candidate = Mock() @@ -247,6 +253,7 @@ async def test_run_prompt_async_returns_text(mock_client): client = Mock(Client) models = Mock(Models) async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() response = Mock() response.candidates = ["yes!"] @@ -274,6 +281,7 @@ async def test_run_prompt_async_raises_when_no_output(mock_client): client = Mock(Client) models = Mock(Models) async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() response = Mock() response.candidates = None @@ -293,3 +301,64 @@ async def get_response(): "do something", model_config=ModelConfig(project="p", location="l", model_name="model"), ) + + +# --- client lifecycle (no per-call leak) --- + + +@patch("genai_utils.gemini.genai.Client") +def test_run_prompt_reuses_cached_client(mock_client): + """The sync path builds one client per project/location and reuses it, + rather than constructing (and never closing) a fresh client per call.""" + _get_client.cache_clear() + try: + client = Mock(Client) + models = Mock(Models) + + response = Mock() + response.candidates = ["yes!"] + response.text = "response!" + models.generate_content.return_value = response + + client.models = models + mock_client.return_value = client + + config = ModelConfig(project="p", location="l", model_name="gemini-2.0-flash") + assert run_prompt("first", model_config=config) == "response!" + assert run_prompt("second", model_config=config) == "response!" + + # Constructed once and reused for both calls. + assert mock_client.call_count == 1 + assert models.generate_content.call_count == 2 + finally: + _get_client.cache_clear() + + +@patch("genai_utils.gemini.genai.Client") +async def test_run_prompt_async_closes_client(mock_client): + """The async path closes the client's async transport so its connection + pool isn't leaked: the loop it binds to is often short-lived, so the client + can't be cached and reused the way the sync one is.""" + client = Mock(Client) + models = Mock(Models) + async_client = Mock(AsyncClient) + async_client.aclose = AsyncMock() + + response = Mock() + response.candidates = ["yes!"] + response.text = "response!" + + async def get_response(): + return response + + models.generate_content.return_value = get_response() + client.aio = async_client + async_client.models = models + mock_client.return_value = client + + result = await run_prompt_async( + "do something", + model_config=ModelConfig(project="p", location="l", model_name="model"), + ) + assert result == "response!" + async_client.aclose.assert_awaited_once() From 02683935dc5c697ed55a7b37a78cd17377392cae Mon Sep 17 00:00:00 2001 From: James McMinn Date: Wed, 17 Jun 2026 15:41:29 +0100 Subject: [PATCH 2/2] fix: don't use a cache for sync calls --- src/genai_utils/gemini.py | 26 +++++++++-------------- tests/genai_utils/test_gemini.py | 36 +++++++++++++------------------- 2 files changed, 24 insertions(+), 38 deletions(-) diff --git a/src/genai_utils/gemini.py b/src/genai_utils/gemini.py index 552dadb..53bbe97 100644 --- a/src/genai_utils/gemini.py +++ b/src/genai_utils/gemini.py @@ -1,7 +1,6 @@ import logging import os import re -from functools import lru_cache from typing import Any import requests @@ -343,19 +342,6 @@ def get_thinking_config( return None -@lru_cache(maxsize=None) -def _get_client(project: str, location: str) -> genai.Client: - """Return a process-wide Vertex AI client for a project/location pair. - - The client owns HTTP connection pools, so it is built once and reused for - every call rather than constructed per request. A fresh client each call is - never closed, so its pools accumulate and memory grows steadily under load. - google-genai's sync client is safe to share across threads, and the set of - distinct (project, location) pairs is tiny, so caching every pair for the - process lifetime is bounded.""" - return genai.Client(vertexai=True, project=project, location=location) - - def _build_request( prompt: str, video_uri: str | None, @@ -535,8 +521,16 @@ class Movie(BaseModel): inline_citations, labels, ) - client = _get_client(model_config.project, model_config.location) - response = client.models.generate_content(**request) + client = genai.Client( + vertexai=True, + project=model_config.project, + location=model_config.location, + ) + try: + response = client.models.generate_content(**request) + finally: + client.close() + return _parse_response(response, use_grounding, inline_citations) diff --git a/tests/genai_utils/test_gemini.py b/tests/genai_utils/test_gemini.py index aff083b..99e0c1d 100644 --- a/tests/genai_utils/test_gemini.py +++ b/tests/genai_utils/test_gemini.py @@ -12,7 +12,6 @@ GeminiError, ModelConfig, NoGroundingError, - _get_client, generate_model_config, get_thinking_config, run_prompt, @@ -307,31 +306,24 @@ async def get_response(): @patch("genai_utils.gemini.genai.Client") -def test_run_prompt_reuses_cached_client(mock_client): - """The sync path builds one client per project/location and reuses it, - rather than constructing (and never closing) a fresh client per call.""" - _get_client.cache_clear() - try: - client = Mock(Client) - models = Mock(Models) +def test_run_prompt_closes_client(mock_client): + """The sync path closes the client after each call so its connection pool + isn't leaked, mirroring the async path.""" + client = Mock(Client) + models = Mock(Models) - response = Mock() - response.candidates = ["yes!"] - response.text = "response!" - models.generate_content.return_value = response + response = Mock() + response.candidates = ["yes!"] + response.text = "response!" + models.generate_content.return_value = response - client.models = models - mock_client.return_value = client + client.models = models + mock_client.return_value = client - config = ModelConfig(project="p", location="l", model_name="gemini-2.0-flash") - assert run_prompt("first", model_config=config) == "response!" - assert run_prompt("second", model_config=config) == "response!" + config = ModelConfig(project="p", location="l", model_name="gemini-2.0-flash") + assert run_prompt("do something", model_config=config) == "response!" - # Constructed once and reused for both calls. - assert mock_client.call_count == 1 - assert models.generate_content.call_count == 2 - finally: - _get_client.cache_clear() + client.close.assert_called_once() @patch("genai_utils.gemini.genai.Client")