fix: reuse Gemini client instead of building one per call

src/genai_utils/gemini.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 import os
 import re
@@ -343,6 +342,88 @@ def get_thinking_config(
     return None
 
 
+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 +505,33 @@ 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 = 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)
 
 
 async def run_prompt_async(
@@ -522,69 +615,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)

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
@@ -14,6 +14,7 @@
     NoGroundingError,
     generate_model_config,
     get_thinking_config,
+    run_prompt,
     run_prompt_async,
     validate_labels,
 )
@@ -67,6 +68,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 +102,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 +130,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 +158,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 +252,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 +280,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 +300,57 @@ 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_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
+
+    client.models = models
+    mock_client.return_value = client
+
+    config = ModelConfig(project="p", location="l", model_name="gemini-2.0-flash")
+    assert run_prompt("do something", model_config=config) == "response!"
+
+    client.close.assert_called_once()
+
+
+@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()