From 409b0fd75a099ea6cd9a65d8f837d2af4ab0e3c4 Mon Sep 17 00:00:00 2001 From: Ed Dearden Date: Fri, 27 Feb 2026 11:05:41 +0000 Subject: [PATCH 1/3] feat: raise an exception when grounding doesn't run --- src/genai_utils/gemini.py | 21 +++++++++++++++------ tests/genai_utils/test_gemini.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/src/genai_utils/gemini.py b/src/genai_utils/gemini.py index 63d41e1..d27141f 100644 --- a/src/genai_utils/gemini.py +++ b/src/genai_utils/gemini.py @@ -45,12 +45,18 @@ } -class GeminiError(Exception): +class GeminiError(RuntimeError): """ Exception raised when something goes wrong with Gemini. """ +class NoGroundingError(GeminiError): + """ + Exception raised if grounding doesn't run when asked. + """ + + class ModelConfig(BaseModel): """ Config for a Gemini model. @@ -564,18 +570,21 @@ class Movie(BaseModel): ), ) + 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.warning( + _logger.error( "Grounding Info: GROUNDING FAILED - see previous log messages for reason" ) + raise NoGroundingError("Grounding did not run") - if response.candidates and response.text and isinstance(response.text, str): if inline_citations and response.candidates[0].grounding_metadata: text_with_citations = add_citations(response) return text_with_citations - else: - return response.text - raise GeminiError(f"No model output: possible reason: {response.prompt_feedback}") + return response.text diff --git a/tests/genai_utils/test_gemini.py b/tests/genai_utils/test_gemini.py index 82d4be0..dd2eb45 100644 --- a/tests/genai_utils/test_gemini.py +++ b/tests/genai_utils/test_gemini.py @@ -5,12 +5,13 @@ from google.genai.client import AsyncClient from google.genai.models import Models from pydantic import BaseModel, Field -from pytest import mark, param +from pytest import mark, param, raises from genai_utils.gemini import ( DEFAULT_PARAMETERS, GeminiError, ModelConfig, + NoGroundingError, generate_model_config, get_thinking_config, run_prompt_async, @@ -147,6 +148,35 @@ async def test_error_if_citations_and_no_grounding(mock_client): assert False +@patch("genai_utils.gemini.genai.Client") +async def test_no_grounding_error_when_grounding_does_not_run(mock_client): + client = Mock(Client) + models = Mock(Models) + async_client = Mock(AsyncClient) + + async def get_no_grounding_metadata_response(): + candidate = Mock() + candidate.grounding_metadata = None + response = Mock() + response.candidates = [candidate] + response.text = "response!" + return response + + models.generate_content.return_value = get_no_grounding_metadata_response() + client.aio = async_client + async_client.models = models + mock_client.return_value = client + + with raises(NoGroundingError): + await run_prompt_async( + "do something", + use_grounding=True, + model_config=ModelConfig( + project="project", location="location", model_name="model" + ), + ) + + @mark.parametrize( "model_name,do_thinking,expected", [ From 28643f0459a29095a02760f4c2ba0e708358192b Mon Sep 17 00:00:00 2001 From: Ed Dearden Date: Fri, 27 Feb 2026 11:23:16 +0000 Subject: [PATCH 2/3] tests: increased test coverage --- tests/genai_utils/test_gemini.py | 225 +++++++++++++++++++++++++++++++ 1 file changed, 225 insertions(+) diff --git a/tests/genai_utils/test_gemini.py b/tests/genai_utils/test_gemini.py index dd2eb45..5d1e654 100644 --- a/tests/genai_utils/test_gemini.py +++ b/tests/genai_utils/test_gemini.py @@ -1,6 +1,7 @@ import os from unittest.mock import Mock, patch +import requests as req from google.genai import Client, types from google.genai.client import AsyncClient from google.genai.models import Models @@ -12,9 +13,14 @@ GeminiError, ModelConfig, NoGroundingError, + add_citations, + check_grounding_ran, + follow_redirect, generate_model_config, get_thinking_config, + insert_citation, run_prompt_async, + validate_labels, ) @@ -205,3 +211,222 @@ def test_get_thinking_config( ): thinking_config = get_thinking_config(model_name, do_thinking) assert thinking_config == expected + + +# --- follow_redirect --- + + +@patch("genai_utils.gemini.requests.get") +def test_follow_redirect_success(mock_get): + mock_response = Mock() + mock_response.url = "https://real-url.com" + mock_get.return_value = mock_response + + result = follow_redirect("https://short.url") + assert result == "https://real-url.com" + + +@mark.parametrize( + "exception", + [ + param(req.exceptions.HTTPError("error"), id="http-error"), + param(Exception("something went wrong"), id="generic-exception"), + ], +) +@patch("genai_utils.gemini.requests.get") +def test_follow_redirect_error_falls_back(mock_get, exception): + mock_get.side_effect = exception + assert follow_redirect("https://original.url") == "https://original.url" + + +# --- insert_citation --- + + +@mark.parametrize( + "text,citation_idx,expected", + [ + param("Hello world.", 12, "Hello world. (cite)", id="sentence-end"), + param("Hello world", 5, "Hello (cite) world", id="whitespace-fallback"), + param("Helloworld", 5, "Helloworld (cite)", id="no-whitespace-append"), + ], +) +def test_insert_citation(text, citation_idx, expected): + assert insert_citation(text, "(cite)", citation_idx) == expected + + +# --- validate_labels --- + + +def test_validate_labels_valid(): + labels = {"valid-key": "valid-value", "another_key": "value-123"} + assert validate_labels(labels) == labels + + +@mark.parametrize( + "labels", + [ + param({"": "value"}, id="empty-key"), + param({"a" * 64: "value"}, id="key-too-long"), + param({"key": "a" * 64}, id="value-too-long"), + param({"1key": "value"}, id="key-starts-with-digit"), + param({"_key": "value"}, id="key-starts-with-underscore"), + param({"KEY": "value"}, id="key-uppercase"), + param({"key.dots": "value"}, id="key-with-dots"), + param({"key": "VALUE"}, id="value-uppercase"), + param({"key": "val ue"}, id="value-with-space"), + ], +) +def test_validate_labels_invalid_input_dropped(labels): + assert validate_labels(labels) == {} + + +def test_validate_labels_mixed_keeps_only_valid(): + labels = {"valid": "ok", "INVALID": "value", "": "empty"} + assert validate_labels(labels) == {"valid": "ok"} + + +# --- check_grounding_ran --- + + +def test_check_grounding_ran_no_candidates(): + response = Mock() + response.candidates = [] + assert check_grounding_ran(response) is False + + +def test_check_grounding_ran_no_grounding_metadata(): + candidate = Mock() + candidate.grounding_metadata = None + response = Mock() + response.candidates = [candidate] + assert check_grounding_ran(response) is False + + +def test_check_grounding_ran_returns_true_when_grounding_present(): + metadata = Mock() + metadata.web_search_queries = ["query"] + metadata.grounding_chunks = [Mock()] + metadata.grounding_supports = [Mock()] + candidate = Mock() + candidate.grounding_metadata = metadata + response = Mock() + response.candidates = [candidate] + assert check_grounding_ran(response) is True + + +def test_check_grounding_ran_returns_false_when_no_searches(): + metadata = Mock() + metadata.web_search_queries = [] + metadata.grounding_chunks = [Mock()] + metadata.grounding_supports = [Mock()] + candidate = Mock() + candidate.grounding_metadata = metadata + response = Mock() + response.candidates = [candidate] + assert check_grounding_ran(response) is False + + +# --- add_citations --- + + +@mark.parametrize( + "candidates,text", + [ + param(None, None, id="no-candidates"), + param([Mock()], None, id="no-text"), + ], +) +def test_add_citations_raises_when_missing_output(candidates, text): + response = Mock() + response.candidates = candidates + response.text = text + response.prompt_feedback = "blocked" + with raises(GeminiError): + add_citations(response) + + +def test_add_citations_returns_plain_text_when_no_grounding_metadata(): + candidate = Mock() + candidate.grounding_metadata = None + response = Mock() + response.candidates = [candidate] + response.text = "plain text" + assert add_citations(response) == "plain text" + + +def test_add_citations_returns_plain_text_when_no_supports(): + metadata = Mock() + metadata.grounding_supports = None + metadata.grounding_chunks = [Mock()] + candidate = Mock() + candidate.grounding_metadata = metadata + response = Mock() + response.candidates = [candidate] + response.text = "plain text" + assert add_citations(response) == "plain text" + + +def test_add_citations_returns_plain_text_when_no_chunks(): + metadata = Mock() + metadata.grounding_supports = [Mock()] + metadata.grounding_chunks = None + candidate = Mock() + candidate.grounding_metadata = metadata + response = Mock() + response.candidates = [candidate] + response.text = "plain text" + assert add_citations(response) == "plain text" + + +# --- run_prompt_async happy path --- + + +@patch("genai_utils.gemini.genai.Client") +async def test_run_prompt_async_returns_text(mock_client): + client = Mock(Client) + models = Mock(Models) + async_client = Mock(AsyncClient) + + 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="gemini-2.0-flash"), + ) + assert result == "response!" + + +@patch("genai_utils.gemini.genai.Client") +async def test_run_prompt_async_raises_when_no_output(mock_client): + client = Mock(Client) + models = Mock(Models) + async_client = Mock(AsyncClient) + + response = Mock() + response.candidates = None + response.text = None + response.prompt_feedback = "blocked" + + 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 + + with raises(GeminiError): + await run_prompt_async( + "do something", + model_config=ModelConfig(project="p", location="l", model_name="model"), + ) From 19ebe01b22136c6f47e5aa76d9d78d88283fd846 Mon Sep 17 00:00:00 2001 From: Ed Dearden Date: Fri, 27 Feb 2026 11:39:22 +0000 Subject: [PATCH 3/3] tests: remove duplicate tests from test_gemini, that were already in test_grounding --- tests/genai_utils/test_gemini.py | 143 +--------------------------- tests/genai_utils/test_grounding.py | 25 ++++- 2 files changed, 27 insertions(+), 141 deletions(-) diff --git a/tests/genai_utils/test_gemini.py b/tests/genai_utils/test_gemini.py index 5d1e654..1df6e07 100644 --- a/tests/genai_utils/test_gemini.py +++ b/tests/genai_utils/test_gemini.py @@ -1,7 +1,6 @@ import os from unittest.mock import Mock, patch -import requests as req from google.genai import Client, types from google.genai.client import AsyncClient from google.genai.models import Models @@ -13,12 +12,8 @@ GeminiError, ModelConfig, NoGroundingError, - add_citations, - check_grounding_ran, - follow_redirect, generate_model_config, get_thinking_config, - insert_citation, run_prompt_async, validate_labels, ) @@ -213,47 +208,6 @@ def test_get_thinking_config( assert thinking_config == expected -# --- follow_redirect --- - - -@patch("genai_utils.gemini.requests.get") -def test_follow_redirect_success(mock_get): - mock_response = Mock() - mock_response.url = "https://real-url.com" - mock_get.return_value = mock_response - - result = follow_redirect("https://short.url") - assert result == "https://real-url.com" - - -@mark.parametrize( - "exception", - [ - param(req.exceptions.HTTPError("error"), id="http-error"), - param(Exception("something went wrong"), id="generic-exception"), - ], -) -@patch("genai_utils.gemini.requests.get") -def test_follow_redirect_error_falls_back(mock_get, exception): - mock_get.side_effect = exception - assert follow_redirect("https://original.url") == "https://original.url" - - -# --- insert_citation --- - - -@mark.parametrize( - "text,citation_idx,expected", - [ - param("Hello world.", 12, "Hello world. (cite)", id="sentence-end"), - param("Hello world", 5, "Hello (cite) world", id="whitespace-fallback"), - param("Helloworld", 5, "Helloworld (cite)", id="no-whitespace-append"), - ], -) -def test_insert_citation(text, citation_idx, expected): - assert insert_citation(text, "(cite)", citation_idx) == expected - - # --- validate_labels --- @@ -285,99 +239,6 @@ def test_validate_labels_mixed_keeps_only_valid(): assert validate_labels(labels) == {"valid": "ok"} -# --- check_grounding_ran --- - - -def test_check_grounding_ran_no_candidates(): - response = Mock() - response.candidates = [] - assert check_grounding_ran(response) is False - - -def test_check_grounding_ran_no_grounding_metadata(): - candidate = Mock() - candidate.grounding_metadata = None - response = Mock() - response.candidates = [candidate] - assert check_grounding_ran(response) is False - - -def test_check_grounding_ran_returns_true_when_grounding_present(): - metadata = Mock() - metadata.web_search_queries = ["query"] - metadata.grounding_chunks = [Mock()] - metadata.grounding_supports = [Mock()] - candidate = Mock() - candidate.grounding_metadata = metadata - response = Mock() - response.candidates = [candidate] - assert check_grounding_ran(response) is True - - -def test_check_grounding_ran_returns_false_when_no_searches(): - metadata = Mock() - metadata.web_search_queries = [] - metadata.grounding_chunks = [Mock()] - metadata.grounding_supports = [Mock()] - candidate = Mock() - candidate.grounding_metadata = metadata - response = Mock() - response.candidates = [candidate] - assert check_grounding_ran(response) is False - - -# --- add_citations --- - - -@mark.parametrize( - "candidates,text", - [ - param(None, None, id="no-candidates"), - param([Mock()], None, id="no-text"), - ], -) -def test_add_citations_raises_when_missing_output(candidates, text): - response = Mock() - response.candidates = candidates - response.text = text - response.prompt_feedback = "blocked" - with raises(GeminiError): - add_citations(response) - - -def test_add_citations_returns_plain_text_when_no_grounding_metadata(): - candidate = Mock() - candidate.grounding_metadata = None - response = Mock() - response.candidates = [candidate] - response.text = "plain text" - assert add_citations(response) == "plain text" - - -def test_add_citations_returns_plain_text_when_no_supports(): - metadata = Mock() - metadata.grounding_supports = None - metadata.grounding_chunks = [Mock()] - candidate = Mock() - candidate.grounding_metadata = metadata - response = Mock() - response.candidates = [candidate] - response.text = "plain text" - assert add_citations(response) == "plain text" - - -def test_add_citations_returns_plain_text_when_no_chunks(): - metadata = Mock() - metadata.grounding_supports = [Mock()] - metadata.grounding_chunks = None - candidate = Mock() - candidate.grounding_metadata = metadata - response = Mock() - response.candidates = [candidate] - response.text = "plain text" - assert add_citations(response) == "plain text" - - # --- run_prompt_async happy path --- @@ -401,7 +262,9 @@ async def get_response(): result = await run_prompt_async( "do something", - model_config=ModelConfig(project="p", location="l", model_name="gemini-2.0-flash"), + model_config=ModelConfig( + project="p", location="l", model_name="gemini-2.0-flash" + ), ) assert result == "response!" diff --git a/tests/genai_utils/test_grounding.py b/tests/genai_utils/test_grounding.py index 10ae6f7..c0da80b 100644 --- a/tests/genai_utils/test_grounding.py +++ b/tests/genai_utils/test_grounding.py @@ -2,9 +2,10 @@ import requests from google.genai import types -from pytest import mark, param +from pytest import mark, param, raises from genai_utils.gemini import ( + GeminiError, add_citations, check_grounding_ran, follow_redirect, @@ -177,6 +178,12 @@ def test_insert_citation( "response, expected", [ param(types.GenerateContentResponse(candidates=None), False), + param( + types.GenerateContentResponse( + candidates=[types.Candidate(grounding_metadata=None)] + ), + False, + ), param( types.GenerateContentResponse( candidates=[types.Candidate(grounding_metadata=dummy_grounding)] @@ -194,3 +201,19 @@ def test_insert_citation( def test_check_grounding_ran(response: types.GenerateContentResponse, expected: bool): did_grounding = check_grounding_ran(response) assert did_grounding == expected + + +@mark.parametrize( + "candidates,text", + [ + param(None, None, id="no-candidates"), + param([Mock()], None, id="no-text"), + ], +) +def test_add_citations_raises_when_missing_output(candidates, text): + response = Mock(types.GenerateContentResponse) + response.candidates = candidates + response.text = text + response.prompt_feedback = "blocked" + with raises(GeminiError): + add_citations(response)