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):