tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 06:22:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(gemini): track web search grounding cost on image generation

Browse Source 
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.
This commit is contained in:
mateo-berri
2026-06-12 05:42:32 +00:00
parent e1a4e79bda
commit 48bdcba2cd
10 changed files with 359 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/llm_cost_calc/utils.py
+37
View File
@@ -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:
litellm/llms/gemini/common_utils.py
+20
View File
@@ -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],
litellm/llms/gemini/image_generation/cost_calculator.py
+20 -16
View File
@@ -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
litellm/llms/gemini/image_generation/transformation.py
+6
View File
@@ -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", [])
litellm/llms/vertex_ai/image_generation/cost_calculator.py
+10 -5
View File
@@ -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
litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
+8 -1
View File
@@ -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
tests/test_litellm/llms/gemini/test_cost_calculator.py
+54
View File
@@ -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
tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py
+87
View File
@@ -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
tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py
+72
View File
@@ -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
tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
+45
View File
@@ -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):