diff --git a/src/genai_utils/gemini.py b/src/genai_utils/gemini.py index d27141f..53bbe97 100644 --- a/src/genai_utils/gemini.py +++ b/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) diff --git a/tests/genai_utils/test_gemini.py b/tests/genai_utils/test_gemini.py index 1df6e07..99e0c1d 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 @@ -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()