From 48bdcba2cd48246237a8ff34ce9b6473b88b4aa8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:42:32 +0000 Subject: [PATCH] fix(gemini): track web search grounding cost on image generation Forwarding Google Search grounding to Gemini and Vertex image generation previously incurred billable grounding charges that never reached LiteLLM spend tracking, because the image cost path returns through the Gemini/Vertex image calculators before built-in tool spend is added. Carry the grounding request count from the response onto the image usage object and bill it with the same per-request web search accounting used for chat completions. --- .../litellm_core_utils/llm_cost_calc/utils.py | 37 ++++++++ litellm/llms/gemini/common_utils.py | 20 +++++ .../image_generation/cost_calculator.py | 36 ++++---- .../gemini/image_generation/transformation.py | 6 ++ .../image_generation/cost_calculator.py | 15 ++-- .../vertex_gemini_transformation.py | 9 +- .../llms/gemini/test_cost_calculator.py | 54 ++++++++++++ ..._gemini_image_generation_transformation.py | 87 +++++++++++++++++++ ...tex_ai_image_generation_cost_calculator.py | 72 +++++++++++++++ ...rtex_ai_image_generation_transformation.py | 45 ++++++++++ 10 files changed, 359 insertions(+), 22 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f39c942f90..f96d8b1134 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -929,6 +929,43 @@ def calculate_image_response_cost_from_usage( return prompt_cost + completion_cost +def calculate_image_response_web_search_cost( + image_response: ImageResponse, + custom_llm_provider: str, + model_info: ModelInfo, +) -> float: + """ + Cost of Google Search grounding performed during image generation. + + The grounding request count is carried on the image usage object by the + provider transformers; it is billed with the same per-request accounting + used for chat completions. + """ + usage = image_response.usage + if usage is None: + return 0.0 + + web_search_requests = getattr(usage, "web_search_requests", None) + if not web_search_requests: + return 0.0 + + from litellm.llms import get_cost_for_web_search_request + + synthetic_usage = Usage( + prompt_tokens_details=PromptTokensDetailsWrapper( + web_search_requests=web_search_requests + ) + ) + return ( + get_cost_for_web_search_request( + custom_llm_provider=custom_llm_provider, + usage=synthetic_usage, + model_info=model_info, + ) + or 0.0 + ) + + class CostCalculatorUtils: @staticmethod def _call_type_has_image_response(call_type: str) -> bool: diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index b66441c065..4cca2e2b85 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -258,6 +258,26 @@ def map_gemini_image_tools_params( return result +def get_gemini_image_web_search_requests( + response_data: Dict[str, Any], +) -> Optional[int]: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + grounding_metadata: List[Dict[str, Any]] = [] + for candidate in response_data.get("candidates", []): + if not isinstance(candidate, dict): + continue + candidate_grounding = candidate.get("groundingMetadata") + if isinstance(candidate_grounding, list): + grounding_metadata.extend(candidate_grounding) + elif isinstance(candidate_grounding, dict): + grounding_metadata.append(candidate_grounding) + + return VertexGeminiConfig._calculate_web_search_requests(grounding_metadata) + + def get_gemini_image_generation_config( model: str, optional_params: Dict[str, Any], diff --git a/litellm/llms/gemini/image_generation/cost_calculator.py b/litellm/llms/gemini/image_generation/cost_calculator.py index 3c8e69374a..380e2c21e9 100644 --- a/litellm/llms/gemini/image_generation/cost_calculator.py +++ b/litellm/llms/gemini/image_generation/cost_calculator.py @@ -7,6 +7,7 @@ from typing import Any import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_image_response_cost_from_usage, + calculate_image_response_web_search_cost, ) from litellm.types.utils import ImageResponse @@ -23,22 +24,25 @@ def cost_calculator( custom_llm_provider="gemini", ) - if isinstance(image_response, ImageResponse): - token_based_cost = calculate_image_response_cost_from_usage( - model=model, - image_response=image_response, - custom_llm_provider="gemini", - ) - if token_based_cost is not None: - return token_based_cost - - output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 - if isinstance(image_response, ImageResponse): - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images - else: + if not isinstance(image_response, ImageResponse): raise ValueError( f"image_response must be of type ImageResponse got type={type(image_response)}" ) + + web_search_cost = calculate_image_response_web_search_cost( + image_response=image_response, + custom_llm_provider="gemini", + model_info=_model_info, + ) + + token_based_cost = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider="gemini", + ) + if token_based_cost is not None: + return token_based_cost + web_search_cost + + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 + num_images: int = len(image_response.data) if image_response.data else 0 + return output_cost_per_image * num_images + web_search_cost diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 6b190f538e..ebfb0d6883 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -7,6 +7,7 @@ from litellm.llms.base_llm.image_generation.transformation import ( ) from litellm.llms.gemini.common_utils import ( get_gemini_image_generation_config, + get_gemini_image_web_search_requests, is_gemini_image_model, map_gemini_image_tools_params, map_openai_image_params_to_gemini, @@ -227,6 +228,11 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): model_response.usage = transform_gemini_image_usage( response_data["usageMetadata"] ) + web_search_requests = get_gemini_image_web_search_requests(response_data) + if web_search_requests and model_response.usage is not None: + setattr( + model_response.usage, "web_search_requests", web_search_requests + ) else: # Original Imagen format - predictions with generated images predictions = response_data.get("predictions", []) diff --git a/litellm/llms/vertex_ai/image_generation/cost_calculator.py b/litellm/llms/vertex_ai/image_generation/cost_calculator.py index 012de5498c..5c04ebf79e 100644 --- a/litellm/llms/vertex_ai/image_generation/cost_calculator.py +++ b/litellm/llms/vertex_ai/image_generation/cost_calculator.py @@ -5,6 +5,7 @@ Vertex AI Image Generation Cost Calculator import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_image_response_cost_from_usage, + calculate_image_response_web_search_cost, ) from litellm.types.utils import ImageResponse @@ -21,16 +22,20 @@ def cost_calculator( custom_llm_provider="vertex_ai", ) + web_search_cost = calculate_image_response_web_search_cost( + image_response=image_response, + custom_llm_provider="vertex_ai", + model_info=_model_info, + ) + token_based_cost = calculate_image_response_cost_from_usage( model=model, image_response=image_response, custom_llm_provider="vertex_ai", ) if token_based_cost is not None: - return token_based_cost + return token_based_cost + web_search_cost output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images + num_images: int = len(image_response.data) if image_response.data else 0 + return output_cost_per_image * num_images + web_search_cost diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 2be32f8bea..103c7b2a28 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -7,7 +7,10 @@ import litellm from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from litellm.llms.gemini.common_utils import map_gemini_image_tools_params +from litellm.llms.gemini.common_utils import ( + get_gemini_image_web_search_requests, + map_gemini_image_tools_params, +) from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str @@ -333,4 +336,8 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): if usage_metadata := response_data.get("usageMetadata", None): model_response.usage = self._transform_image_usage(usage_metadata) + web_search_requests = get_gemini_image_web_search_requests(response_data) + if web_search_requests and model_response.usage is not None: + setattr(model_response.usage, "web_search_requests", web_search_requests) + return model_response diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 6d51bcd2c8..6917092966 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -247,3 +247,57 @@ def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(): ) assert cost == len(image_response.data or []) * model_info["output_cost_per_image"] + + +def _image_response_with_web_search(web_search_requests): + usage = ImageUsage( + input_tokens=20, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=20, + image_tokens=0, + ), + output_tokens=1120, + total_tokens=1140, + ) + if web_search_requests is not None: + usage.web_search_requests = web_search_requests + return ImageResponse(data=[ImageObject(b64_json="img1")], usage=usage) + + +def test_gemini_image_generation_cost_adds_web_search_grounding(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + grounded = gemini_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(2), + ) + ungrounded = gemini_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(None), + ) + + expected_web_search_cost = cost_per_web_search_request( + usage=_make_usage(2), model_info=model_info + ) + assert expected_web_search_cost > 0 + assert round(grounded - ungrounded, 10) == round(expected_web_search_cost, 10) + + +def test_gemini_image_generation_cost_no_web_search_when_absent(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini/gemini-3-pro-image-preview" + + cost_zero = gemini_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(0), + ) + cost_none = gemini_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(None), + ) + + assert cost_zero == cost_none diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py index 3a480b0a63..d950985675 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py @@ -332,3 +332,90 @@ def test_gemini_image_generation_usage_without_output_details_treats_output_as_i usage = result.model_dump()["usage"] assert usage["completion_tokens_details"]["text_tokens"] == 0 assert usage["completion_tokens_details"]["image_tokens"] == 1716 + + +def test_gemini_image_generation_response_tracks_web_search_requests(): + config = GoogleImageGenConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "fake-image", + } + } + ] + }, + "groundingMetadata": { + "webSearchQueries": ["latest iphone", "iphone colors"] + }, + } + ], + "usageMetadata": { + "promptTokenCount": 35, + "candidatesTokenCount": 1716, + "totalTokenCount": 1751, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 35}], + }, + }, + ) + + result = config.transform_image_generation_response( + model="gemini-3.1-flash-image-preview", + raw_response=raw_response, + model_response=ImageResponse(data=[]), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.usage.web_search_requests == 2 + + +def test_gemini_image_generation_response_without_grounding_has_no_web_search_requests(): + config = GoogleImageGenConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "fake-image", + } + } + ] + } + } + ], + "usageMetadata": { + "promptTokenCount": 35, + "candidatesTokenCount": 1716, + "totalTokenCount": 1751, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 35}], + }, + }, + ) + + result = config.transform_image_generation_response( + model="gemini-3.1-flash-image-preview", + raw_response=raw_response, + model_response=ImageResponse(data=[]), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert getattr(result.usage, "web_search_requests", None) is None diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py new file mode 100644 index 0000000000..cd86618716 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py @@ -0,0 +1,72 @@ +import os + +import litellm +from litellm.llms.vertex_ai.gemini.cost_calculator import cost_per_web_search_request +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + PromptTokensDetailsWrapper, + Usage, +) + + +def _image_response_with_web_search(web_search_requests): + usage = ImageUsage( + input_tokens=20, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=20, + image_tokens=0, + ), + output_tokens=1120, + total_tokens=1140, + ) + if web_search_requests is not None: + usage.web_search_requests = web_search_requests + return ImageResponse(data=[ImageObject(b64_json="img1")], usage=usage) + + +def test_vertex_image_generation_cost_adds_web_search_grounding(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + + grounded = vertex_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(3), + ) + ungrounded = vertex_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(None), + ) + + expected_web_search_cost = cost_per_web_search_request( + usage=Usage( + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=3) + ), + model_info=model_info, + ) + assert expected_web_search_cost > 0 + assert round(grounded - ungrounded, 10) == round(expected_web_search_cost, 10) + + +def test_vertex_image_generation_cost_no_web_search_when_absent(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini-3-pro-image-preview" + + cost_zero = vertex_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(0), + ) + cost_none = vertex_image_generation_cost_calculator( + model=model, + image_response=_image_response_with_web_search(None), + ) + + assert cost_zero == cost_none diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 275384afbe..dc2d945c33 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -349,6 +349,51 @@ class TestVertexAIGeminiImageGenerationConfig: == "test_signature_abc123" ) + def test_transform_image_generation_response_tracks_web_search_requests(self): + """Grounding queries are carried onto usage so search spend can be billed""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + } + } + ] + }, + "groundingMetadata": { + "webSearchQueries": ["eiffel tower", "paris skyline"] + }, + } + ], + "usageMetadata": { + "promptTokenCount": 93, + "candidatesTokenCount": 17, + "totalTokenCount": 110, + }, + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.usage.web_search_requests == 2 + class TestVertexAIImagenImageGenerationConfig: def setup_method(self):