From 53e1e9b87b36dabed967b17338948535020b7abc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:41:59 -0800 Subject: [PATCH] fix(vertex_ai): bill gemini image edits by token usage (#51038) * 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 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gemini/image_generation/transformation.py | 2 +- .../llms/gemini/image_usage_transformation.py | 6 +- .../vertex_gemini_transformation.py | 18 ++++- ...est_vertex_ai_image_edit_transformation.py | 40 ++++++++++ ...tex_ai_image_generation_cost_calculator.py | 78 +++++++++++++++++++ 5 files changed, 137 insertions(+), 7 deletions(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 1f986b11af2..6d1e1cf3931 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -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, diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index 5f08d8403af..3ef403f79ee 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -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" ) diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 3633bb68b50..8e9bef97085 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -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] diff --git a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index f36480ae7df..e311313c136 100644 --- a/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/unit/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -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: diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py index e54e25cbd18..65d6e97f8e5 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py @@ -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,