scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 455

fix(vertex_ai): bill gemini image edits by token usage (#51038)

Browse source 
* fix(vertex_ai): bill Gemini image edit token usage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(vertex_ai): validate image edit usage before normalization

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: kenneth <kenneth@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-11-10 14:41:59 -08:00 • committed by GitHub
parent ac0d2fcd68
commit 53e1e9b87b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 137 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/gemini/image_generation/transformation.py
View file

@ -170,7 +170,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
return request_body_obj.model_dump(exclude_none=True)
def _transform_image_usage(self, usage_metadata: dict):
return transform_gemini_image_usage(usage_metadata)
return transform_gemini_image_usage(_JSON_DICT.validate_python(usage_metadata))
def transform_image_generation_response(
self,

6
litellm/llms/gemini/image_usage_transformation.py
View file

@ -1,5 +1,7 @@
from typing import Final
from pydantic import TypeAdapter
from litellm.types.utils import ImageUsage, ImageUsageInputTokensDetails
@ -34,12 +36,12 @@ def _sum_modality_token_details(usage_metadata: dict, *details_keys: str) -> Ima
return tokens_details
def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage:
def transform_gemini_image_usage(usage_metadata: dict[str, object]) -> ImageUsage:
"""
Transform Gemini usageMetadata to ImageUsage format.
"""
input_tokens_details = _sum_modality_token_details(usage_metadata, "promptTokensDetails", "prompt_tokens_details")
output_tokens: Final = usage_metadata.get("candidatesTokenCount", 0)
output_tokens: Final = TypeAdapter(int).validate_python(usage_metadata.get("candidatesTokenCount", 0))
output_tokens_details: Final = _sum_modality_token_details(
usage_metadata, "candidatesTokensDetails", "candidates_tokens_details"
)

18
litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py
View file

@ -6,10 +6,12 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, cast
import httpx
from httpx._types import RequestFiles
from pydantic import ConfigDict, TypeAdapter
import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.secret_managers.main import get_secret_str
@ -20,7 +22,9 @@ from litellm.types.llms.vertex_ai import (
HttpxPartType,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse, OpenAIImage
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
_JSON_DICT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -202,7 +206,6 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM):
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ImageResponse:
model_response: Final = ImageResponse()
try:
response_json: Final = _generate_content_payload(raw_response)
except Exception as exc:
@ -223,8 +226,15 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM):
if (inline_data := part.get("inlineData")) and (b64_json := inline_data.get("data"))
]
model_response.data = cast(list[OpenAIImage], data_list)
return model_response
usage_metadata: Final = _JSON_DICT.validate_python(response_json.get("usageMetadata") or {})
traffic_type: Final = usage_metadata.get("trafficType") if usage_metadata else None
return ImageResponse(
data=data_list,
usage=transform_gemini_image_usage(usage_metadata) if usage_metadata else None,
hidden_params={"provider_specific_fields": {"traffic_type": traffic_type}}
if isinstance(traffic_type, str) and traffic_type
else {},
)
def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[HttpxPartType]:
images: Final[list[FileTypes]] = image if isinstance(image, list) else [image]

40
tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py
View file

@ -24,6 +24,46 @@ from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
from litellm.types.router import GenericLiteLLMParams
@pytest.mark.parametrize("traffic_type", ("ON_DEMAND_PRIORITY", "", None))
def test_vertex_gemini_edit_preserves_usage_and_traffic_type(traffic_type: str | None) -> None:
response: Final = VertexAIGeminiImageEditConfig().transform_image_edit_response(
model="gemini-3.1-flash-image",
raw_response=httpx.Response(
200,
json={
"candidates": [],
"usageMetadata": {
"promptTokenCount": 1131,
"candidatesTokenCount": 0,
"totalTokenCount": 1131,
"promptTokensDetails": [
{"modality": "TEXT", "tokenCount": 7},
{"modality": "IMAGE", "tokenCount": 1124},
],
**({"trafficType": traffic_type} if traffic_type is not None else {}),
},
},
),
logging_obj=Logging(
model="gemini-3.1-flash-image",
messages=[],
stream=False,
call_type="image_edit",
start_time=datetime(2025, 1, 1),
litellm_call_id="edit",
function_id="edit",
),
)
assert response.usage is not None
assert response.usage.input_tokens == 1131
assert response.usage.output_tokens == 0
assert response.usage.input_tokens_details.text_tokens == 7
assert response.usage.input_tokens_details.image_tokens == 1124
assert response.hidden_params.get("provider_specific_fields", {}) == (
{"traffic_type": traffic_type} if traffic_type else {}
)
@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096"))
@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image"))
def test_vertex_gemini_edit_size_reaches_request(size: str, model: str) -> None:

78
tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py
View file

@ -1,7 +1,16 @@
import os
from datetime import datetime
from typing import Final
import httpx
import pytest
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.vertex_ai.gemini.cost_calculator import cost_per_web_search_request
from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
VertexAIGeminiImageEditConfig,
)
from litellm.llms.vertex_ai.image_generation.cost_calculator import (
cost_calculator as vertex_image_generation_cost_calculator,
)
@ -15,6 +24,75 @@ from litellm.types.utils import (
)
@pytest.mark.parametrize("call_type", ("image_edit", "aimage_edit"))
@pytest.mark.parametrize("has_image, has_usage", ((True, True), (False, True), (True, False)))
def test_vertex_image_edit_cost_uses_input_tokens_and_retains_image_fallback(
call_type: str,
has_image: bool,
has_usage: bool,
) -> None:
model: Final = "gemini-3.1-flash-image"
model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
response: Final = VertexAIGeminiImageEditConfig().transform_image_edit_response(
model=model,
raw_response=httpx.Response(
200,
json={
"candidates": [
{
"finishReason": "STOP" if has_image else "IMAGE_RECITATION",
**(
{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}
if has_image
else {}
),
}
],
**(
{
"usageMetadata": {
"promptTokenCount": 1131,
"candidatesTokenCount": 100 if has_image else 0,
"totalTokenCount": 1231 if has_image else 1131,
"promptTokensDetails": [
{"modality": "TEXT", "tokenCount": 7},
{"modality": "IMAGE", "tokenCount": 1124},
],
}
}
if has_usage
else {}
),
},
),
logging_obj=Logging(
model=model,
messages=[],
stream=False,
call_type=call_type,
start_time=datetime(2025, 1, 1),
litellm_call_id="edit",
function_id="edit",
),
)
cost: Final = litellm.completion_cost(
model=model,
custom_llm_provider="vertex_ai",
completion_response=response,
call_type=call_type,
)
expected: Final = (
7 * (model_info["input_cost_per_token"] or 0)
+ 1124 * (model_info.get("input_cost_per_image_token") or model_info["input_cost_per_token"] or 0)
+ (100 if has_image else 0) * (model_info.get("output_cost_per_image_token") or 0)
if has_usage
else model_info["output_cost_per_image"]
)
assert expected is not None and expected > 0
assert cost == pytest.approx(expected)
assert vertex_image_generation_cost_calculator(model=model, image_response=response) == pytest.approx(expected)
def _image_response_with_web_search(web_search_requests):
usage = ImageUsage(
input_tokens=20,