scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

fix(vertex_ai): bill gemini image text, thinking and priority tokens at their own rates

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-10-06 20:17:53 +00:00
parent abd3af422d
commit e26fa5ba69
10 changed files with 456 additions and 68 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/cost_calculator.py
View file

@ -1576,6 +1576,7 @@ def completion_cost(
call_type=call_type,
model_info=_deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id),
vertex_location=vertex_location,
service_tier=service_tier,
)
elif call_type in _VIDEO_CALL_TYPES:
### VIDEO GENERATION COST CALCULATION ###

4
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -1776,6 +1776,7 @@ def calculate_image_response_cost_from_usage(
custom_llm_provider: str,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
service_tier: str | None = None,
) -> float | None:
"""
Calculate image generation cost from usage metadata when available.
@ -1859,6 +1860,7 @@ def calculate_image_response_cost_from_usage(
usage=normalized_usage,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
service_tier=service_tier,
vertex_location=vertex_location,
)
return prompt_cost + completion_cost
@ -1922,6 +1924,7 @@ class CostCalculatorUtils:
call_type: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
service_tier: str | None = None,
) -> float:
"""
Route the image generation cost calculator based on the custom_llm_provider
@ -1964,6 +1967,7 @@ class CostCalculatorUtils:
image_response=completion_response,
model_info=pricing,
vertex_location=vertex_location,
service_tier=service_tier,
)
elif custom_llm_provider == litellm.LlmProviders.BEDROCK.value:
if isinstance(completion_response, ImageResponse):

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

@ -1,65 +1,130 @@
from collections.abc import Mapping
from typing import Final
from pydantic import TypeAdapter
from litellm.types.utils import ImageUsage, ImageUsageInputTokensDetails
_TOKEN_COUNT_ADAPTER: Final = TypeAdapter(int)
_TOKEN_DETAILS_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_TOKEN_DETAILS_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _get_token_count(details: dict) -> int:
def _get_token_count(details: Mapping[str, object]) -> int:
raw_token_count: Final = details.get("tokenCount", details.get("token_count", 0))
return raw_token_count if isinstance(raw_token_count, int) else 0
def _get_modality_token_details(usage_metadata: dict, *details_keys: str) -> list:
for details_key in details_keys:
details = usage_metadata.get(details_key)
if isinstance(details, list):
return details
return []
def _get_usage_token_count(usage_metadata: Mapping[str, object], key: str) -> int:
return _TOKEN_COUNT_ADAPTER.validate_python(usage_metadata.get(key, 0))
def _sum_modality_token_details(usage_metadata: dict, *details_keys: str) -> ImageUsageInputTokensDetails:
tokens_details: Final = ImageUsageInputTokensDetails(
image_tokens=0,
text_tokens=0,
def _get_modality_token_details(
usage_metadata: Mapping[str, object], *details_keys: str
) -> tuple[Mapping[str, object], ...]:
details: Final[list[object]] = next(
(
_TOKEN_DETAILS_LIST_ADAPTER.validate_python(usage_metadata.get(details_key))
for details_key in details_keys
if isinstance(usage_metadata.get(details_key), list)
),
[],
)
return tuple(
_TOKEN_DETAILS_MAPPING_ADAPTER.validate_python(detail) for detail in details if isinstance(detail, Mapping)
)
for details in _get_modality_token_details(usage_metadata, *details_keys):
if isinstance(details, dict):
modality = str(details.get("modality", "")).upper()
token_count = _get_token_count(details)
if modality == "TEXT":
tokens_details.text_tokens += token_count
elif modality == "IMAGE":
tokens_details.image_tokens += token_count
return tokens_details
def _sum_modality_token_details(
token_details: tuple[Mapping[str, object], ...],
) -> ImageUsageInputTokensDetails:
return ImageUsageInputTokensDetails(
image_tokens=sum(
_get_token_count(details)
for details in token_details
if str(details.get("modality", "")).upper() == "IMAGE"
),
text_tokens=sum(
_get_token_count(details) for details in token_details if str(details.get("modality", "")).upper() == "TEXT"
),
)
def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage:
def _output_modality_token_details(
candidate_tokens: int,
reasoning_tokens: int,
candidate_tokens_are_inclusive: bool,
candidate_token_details: tuple[Mapping[str, object], ...],
modality_tokens: ImageUsageInputTokensDetails,
) -> ImageUsageInputTokensDetails:
non_reasoning_candidate_tokens: Final = max(
candidate_tokens - (reasoning_tokens if candidate_tokens_are_inclusive else 0),
0,
)
if not candidate_token_details:
return ImageUsageInputTokensDetails(image_tokens=non_reasoning_candidate_tokens, text_tokens=0)
known_modality_tokens: Final = modality_tokens.text_tokens + modality_tokens.image_tokens
text_tokens: Final = (
max(
modality_tokens.text_tokens - (known_modality_tokens - non_reasoning_candidate_tokens),
0,
)
if known_modality_tokens > non_reasoning_candidate_tokens
else modality_tokens.text_tokens + non_reasoning_candidate_tokens - known_modality_tokens
)
return ImageUsageInputTokensDetails(image_tokens=modality_tokens.image_tokens, text_tokens=text_tokens)
def transform_gemini_image_usage(usage_metadata: Mapping[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_details: Final = _sum_modality_token_details(
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
validated_counts: Final = ImageUsage.model_validate(
{
"input_tokens": usage_metadata.get("promptTokenCount", 0),
"input_tokens_details": {"image_tokens": 0, "text_tokens": 0},
"output_tokens": usage_metadata.get("candidatesTokenCount", 0),
"total_tokens": usage_metadata.get("totalTokenCount", 0),
}
)
prompt_tokens: Final = validated_counts.input_tokens
candidate_tokens: Final = validated_counts.output_tokens
reasoning_tokens: Final = _get_usage_token_count(usage_metadata, "thoughtsTokenCount")
candidate_tokens_are_inclusive: Final = VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata)
completion_tokens: Final = candidate_tokens + (0 if candidate_tokens_are_inclusive else reasoning_tokens)
input_tokens_details: Final = _sum_modality_token_details(
_get_modality_token_details(usage_metadata, "promptTokensDetails", "prompt_tokens_details")
)
candidate_token_details: Final = _get_modality_token_details(
usage_metadata, "candidatesTokensDetails", "candidates_tokens_details"
)
modality_tokens: Final = _sum_modality_token_details(candidate_token_details)
output_modality_tokens: Final = _output_modality_token_details(
candidate_tokens,
reasoning_tokens,
candidate_tokens_are_inclusive,
candidate_token_details,
modality_tokens,
)
if not _get_modality_token_details(usage_metadata, "candidatesTokensDetails", "candidates_tokens_details"):
output_tokens_details.image_tokens = output_tokens
else:
known_output_tokens: Final = output_tokens_details.text_tokens + output_tokens_details.image_tokens
if output_tokens > known_output_tokens:
output_tokens_details.text_tokens += output_tokens - known_output_tokens
output_tokens_details: Final[dict[str, int]] = {
"image_tokens": output_modality_tokens.image_tokens,
"text_tokens": output_modality_tokens.text_tokens,
**({"reasoning_tokens": reasoning_tokens} if "thoughtsTokenCount" in usage_metadata else {}),
}
usage_payload: Final[dict[str, object]] = {
"input_tokens": usage_metadata.get("promptTokenCount", 0),
"input_tokens": prompt_tokens,
"input_tokens_details": input_tokens_details,
"output_tokens": output_tokens,
"total_tokens": usage_metadata.get("totalTokenCount", 0),
"prompt_tokens": usage_metadata.get("promptTokenCount", 0),
"output_tokens": completion_tokens,
"total_tokens": validated_counts.total_tokens,
"prompt_tokens": prompt_tokens,
"prompt_tokens_details": input_tokens_details.model_dump(),
"completion_tokens": output_tokens,
"completion_tokens_details": output_tokens_details.model_dump(),
"output_tokens_details": output_tokens_details.model_dump(),
"completion_tokens": completion_tokens,
"completion_tokens_details": {**output_tokens_details},
"output_tokens_details": {**output_tokens_details},
}
return ImageUsage.model_validate(usage_payload)

18
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -65,7 +65,6 @@ from litellm.types.llms.vertex_ai import (
LogprobsResult,
ToolConfig,
Tools,
UsageMetadata,
VertexToolName,
)
from litellm.types.utils import (
@ -1696,7 +1695,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return model_response
@staticmethod
def is_candidate_token_count_inclusive(usage_metadata: UsageMetadata) -> bool:
def is_candidate_token_count_inclusive(usage_metadata: Mapping[str, object]) -> bool:
"""
Check if the candidate token count is inclusive of the thinking token count
@ -1706,12 +1705,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
Addresses - https://github.com/BerriAI/litellm/pull/10141#discussion_r2052272035
"""
non_thinking_tokens: Final = (
usage_metadata.get("promptTokenCount", 0)
+ usage_metadata.get("candidatesTokenCount", 0)
+ usage_metadata.get("toolUsePromptTokenCount", 0)
prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0)
candidate_tokens: Final = usage_metadata.get("candidatesTokenCount", 0)
tool_use_prompt_tokens: Final = usage_metadata.get("toolUsePromptTokenCount", 0)
total_tokens: Final = usage_metadata.get("totalTokenCount", 0)
return (
isinstance(prompt_tokens, int)
and isinstance(candidate_tokens, int)
and isinstance(tool_use_prompt_tokens, int)
and isinstance(total_tokens, int)
and prompt_tokens + candidate_tokens + tool_use_prompt_tokens == total_tokens
)
return non_thinking_tokens == usage_metadata.get("totalTokenCount", 0)
@staticmethod
def _response_has_search_grounding(

2
litellm/llms/vertex_ai/image_generation/cost_calculator.py
View file

@ -18,6 +18,7 @@ def cost_calculator(
image_response: ImageResponse,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
service_tier: str | None = None,
) -> float:
"""
Vertex AI Image Generation Cost Calculator
@ -40,6 +41,7 @@ def cost_calculator(
custom_llm_provider="vertex_ai",
model_info=_model_info,
vertex_location=vertex_location,
service_tier=service_tier,
)
if token_based_cost is not None:
return token_based_cost + web_search_cost

42
litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
View file

@ -14,6 +14,7 @@ from litellm.llms.gemini.common_utils import (
get_gemini_image_web_search_requests,
map_gemini_image_tools_params,
)
from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage
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
@ -21,8 +22,6 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
ImageObject,
ImageResponse,
ImageUsage,
ImageUsageInputTokensDetails,
)
if TYPE_CHECKING:
@ -258,27 +257,6 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
return request_body
def _transform_image_usage(self, usage: dict) -> ImageUsage:
input_tokens_details: Final = ImageUsageInputTokensDetails(
image_tokens=0,
text_tokens=0,
)
tokens_details: Final = usage.get("promptTokensDetails", [])
for details in tokens_details:
if isinstance(details, dict) and (modality := details.get("modality")):
token_count = details.get("tokenCount", 0)
if modality == "TEXT":
input_tokens_details.text_tokens += token_count
elif modality == "IMAGE":
input_tokens_details.image_tokens += token_count
return ImageUsage(
input_tokens=usage.get("promptTokenCount", 0),
input_tokens_details=input_tokens_details,
output_tokens=usage.get("candidatesTokenCount", 0),
total_tokens=usage.get("totalTokenCount", 0),
)
def transform_image_generation_response(
self,
model: str,
@ -328,7 +306,23 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
response_object: Final = _JSON_OBJECT.validate_python(response_data)
if usage_metadata := response_object.get("usageMetadata", None):
model_response.usage = self._transform_image_usage(_JSON_DICT.validate_python(usage_metadata))
normalized_usage_metadata: Final = _JSON_DICT.validate_python(usage_metadata)
model_response.usage = transform_gemini_image_usage(normalized_usage_metadata)
traffic_type: Final = normalized_usage_metadata.get("trafficType")
if isinstance(traffic_type, str) and traffic_type:
hidden_params: Final[dict[str, object]] = _JSON_DICT.validate_python(model_response._hidden_params) # pyright: ignore[reportPrivateUsage] # ImageResponse has no public hidden-parameter accessor
provider_specific_fields: Final[dict[str, object]] = (
_JSON_DICT.validate_python(hidden_params.get("provider_specific_fields"))
if isinstance(hidden_params.get("provider_specific_fields"), Mapping)
else {}
)
model_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ImageResponse has no public setter
**hidden_params,
"provider_specific_fields": {
**provider_specific_fields,
"traffic_type": traffic_type,
},
}
web_search_requests: Final = get_gemini_image_web_search_requests(response_object)
if web_search_requests and model_response.usage is not None:

134
tests/integration/spend/test_vertex_gemini_image_token_spend.py Normal file
View file

@ -0,0 +1,134 @@
import json
from collections.abc import Callable
from hashlib import sha256
from typing import Final
import pytest
from integration._support.client import Gateway, Scenario, eventually
from integration._support.database import read_rows
from integration._support.vertex import service_account_json
from integration._support.wire import Reply, Request, Wire, wire_server
_PROJECT: Final = "scripted-project"
_LOCATION: Final = "global"
_BACKEND: Final = "gemini-scripted-image-tokens"
_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models"
_PROMPT: Final = "scripted image token spend"
_PNG_B64: Final = (
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
)
_PROMPT_TOKENS: Final = 12
_TEXT_OUTPUT_TOKENS: Final = 10
_IMAGE_OUTPUT_TOKENS: Final = 1290
_THOUGHTS_TOKENS: Final = 40
_INPUT_RATE: Final = 2e-6
_TEXT_OUTPUT_RATE: Final = 1e-5
_IMAGE_OUTPUT_RATE: Final = 4e-5
_PRIORITY_INPUT_RATE: Final = 4e-6
_PRIORITY_OUTPUT_RATE: Final = 2e-5
_EXPECTED_STANDARD_SPEND: Final = (
_PROMPT_TOKENS * _INPUT_RATE
+ (_TEXT_OUTPUT_TOKENS + _THOUGHTS_TOKENS) * _TEXT_OUTPUT_RATE
+ _IMAGE_OUTPUT_TOKENS * _IMAGE_OUTPUT_RATE
)
_EXPECTED_PRIORITY_SPEND: Final = (
_PROMPT_TOKENS * _PRIORITY_INPUT_RATE
+ (_TEXT_OUTPUT_TOKENS + _THOUGHTS_TOKENS) * _PRIORITY_OUTPUT_RATE
+ _IMAGE_OUTPUT_TOKENS * _IMAGE_OUTPUT_RATE
)
def _image_reply(traffic_type: str | None) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert (request.method, request.target) == (
"POST",
f"{_MODEL_PATH}/{_BACKEND}:generateContent",
), (request.method, request.target)
assert _PROMPT in request.body.decode(), request.body
usage_metadata: Final[dict[str, object]] = {
"promptTokenCount": _PROMPT_TOKENS,
"candidatesTokenCount": _TEXT_OUTPUT_TOKENS + _IMAGE_OUTPUT_TOKENS,
"thoughtsTokenCount": _THOUGHTS_TOKENS,
"totalTokenCount": _PROMPT_TOKENS + _TEXT_OUTPUT_TOKENS + _IMAGE_OUTPUT_TOKENS + _THOUGHTS_TOKENS,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": _PROMPT_TOKENS}],
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": _TEXT_OUTPUT_TOKENS},
{"modality": "IMAGE", "tokenCount": _IMAGE_OUTPUT_TOKENS},
],
**({"trafficType": traffic_type} if traffic_type is not None else {}),
}
return Reply(
body=json.dumps(
{
"candidates": [
{
"content": {
"role": "model",
"parts": [
{"inlineData": {"mimeType": "image/png", "data": _PNG_B64}},
],
},
"finishReason": "STOP",
}
],
"usageMetadata": usage_metadata,
}
).encode()
)
return respond
def _deployment(gateway: Gateway, scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model=f"vertex_ai/{_BACKEND}",
api_base=f"{wire.url}{_MODEL_PATH}/{_BACKEND}:generateContent",
api_key=None,
vertex_project=_PROJECT,
vertex_location=_LOCATION,
vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
input_cost_per_token=_INPUT_RATE,
output_cost_per_token=_TEXT_OUTPUT_RATE,
output_cost_per_image_token=_IMAGE_OUTPUT_RATE,
input_cost_per_token_priority=_PRIORITY_INPUT_RATE,
output_cost_per_token_priority=_PRIORITY_OUTPUT_RATE,
)
def _spend_for_key(key: str) -> float:
rows: Final = eventually(
lambda: read_rows(
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(sha256(key.encode()).hexdigest(),),
),
lambda found: len(found) == 1,
seconds=70,
)
return float(str(rows[0]["spend"]))
@pytest.mark.timeout(120)
@pytest.mark.parametrize(
("traffic_type", "expected_spend"),
[
pytest.param(None, _EXPECTED_STANDARD_SPEND, id="standard"),
pytest.param("ON_DEMAND_PRIORITY", _EXPECTED_PRIORITY_SPEND, id="on-demand-priority"),
],
)
def test_vertex_gemini_image_spend_prices_text_thinking_and_served_tier(
gateway: Gateway,
traffic_type: str | None,
expected_spend: float,
) -> None:
with wire_server(_image_reply(traffic_type)) as wire, gateway.scenario() as scenario:
model: Final = _deployment(gateway, scenario, wire)
key: Final = scenario.key(models=[model])
response: Final = gateway.request(
"POST",
"/v1/images/generations",
{"model": model, "prompt": _PROMPT},
key=key,
)
assert response.status_code == 200, response.text
assert _spend_for_key(key) == pytest.approx(expected_spend)

67
tests/unit/llms/gemini/test_image_usage_transformation.py
View file

@ -45,3 +45,70 @@ def test_transform_gemini_image_usage_reports_image_and_chat_style_counts(
"completion_tokens_details": expected_output_details,
"output_tokens_details": expected_output_details,
}
def test_transform_gemini_image_usage_adds_exclusive_thoughts_to_output_tokens() -> None:
usage: Final = transform_gemini_image_usage(
{
"promptTokenCount": 12,
"candidatesTokenCount": 1300,
"thoughtsTokenCount": 40,
"totalTokenCount": 1352,
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": 10},
{"modality": "IMAGE", "tokenCount": 1290},
],
}
)
assert usage.output_tokens == 1340
assert usage.completion_tokens == 1340
assert usage.completion_tokens_details == {
"text_tokens": 10,
"image_tokens": 1290,
"reasoning_tokens": 40,
}
assert usage.output_tokens_details == usage.completion_tokens_details
def test_transform_gemini_image_usage_does_not_double_count_inclusive_thoughts() -> None:
usage: Final = transform_gemini_image_usage(
{
"promptTokenCount": 12,
"candidatesTokenCount": 1300,
"thoughtsTokenCount": 40,
"totalTokenCount": 1312,
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": 10},
{"modality": "IMAGE", "tokenCount": 1250},
],
}
)
assert usage.output_tokens == 1300
assert usage.completion_tokens == 1300
assert usage.completion_tokens_details == {
"text_tokens": 10,
"image_tokens": 1250,
"reasoning_tokens": 40,
}
assert usage.output_tokens_details == usage.completion_tokens_details
def test_transform_gemini_image_usage_fallback_excludes_thoughts_from_image_tokens() -> None:
usage: Final = transform_gemini_image_usage(
{
"promptTokenCount": 12,
"candidatesTokenCount": 1300,
"thoughtsTokenCount": 40,
"totalTokenCount": 1352,
}
)
assert usage.output_tokens == 1340
assert usage.completion_tokens_details == {
"text_tokens": 0,
"image_tokens": 1300,
"reasoning_tokens": 40,
}
assert usage.output_tokens_details == usage.completion_tokens_details

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

@ -1,6 +1,8 @@
import os
from typing import Final
import litellm
import pytest
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,
@ -10,6 +12,7 @@ from litellm.types.utils import (
ImageResponse,
ImageUsage,
ImageUsageInputTokensDetails,
ModelInfo,
PromptTokensDetailsWrapper,
Usage,
)
@ -70,3 +73,61 @@ def test_vertex_image_generation_cost_no_web_search_when_absent(monkeypatch):
)
assert cost_zero == cost_none
def test_vertex_image_generation_cost_prices_token_details_by_service_tier() -> None:
model: Final = "gemini-scripted-image-tokens"
model_info: Final[ModelInfo] = {
"key": model,
"max_tokens": 1000,
"max_input_tokens": 1000,
"max_output_tokens": 1000,
"input_cost_per_token": 2e-6,
"output_cost_per_token": 1e-5,
"output_cost_per_image_token": 4e-5,
"input_cost_per_token_priority": 4e-6,
"output_cost_per_token_priority": 2e-5,
"litellm_provider": "vertex_ai",
"mode": "image_generation",
"supported_openai_params": None,
}
usage: Final = ImageUsage.model_validate(
{
"input_tokens": 12,
"input_tokens_details": {"text_tokens": 12, "image_tokens": 0},
"output_tokens": 1340,
"total_tokens": 1352,
"completion_tokens_details": {
"text_tokens": 10,
"image_tokens": 1290,
"reasoning_tokens": 40,
},
"output_tokens_details": {
"text_tokens": 10,
"image_tokens": 1290,
"reasoning_tokens": 40,
},
}
)
image_response: Final = ImageResponse(
data=[ImageObject(b64_json="synthetic-image")],
usage=usage,
)
standard_cost: Final = vertex_image_generation_cost_calculator(
model=model,
image_response=image_response,
model_info=model_info,
vertex_location="global",
)
priority_cost: Final = vertex_image_generation_cost_calculator(
model=model,
image_response=image_response,
model_info=model_info,
vertex_location="global",
service_tier="priority",
)
expected_standard_cost: Final = 12 * 2e-6 + (10 + 40) * 1e-5 + 1290 * 4e-5
expected_priority_cost: Final = 12 * 4e-6 + (10 + 40) * 2e-5 + 1290 * 4e-5
assert standard_cost == pytest.approx(expected_standard_cost)
assert priority_cost == pytest.approx(expected_priority_cost)

56
tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
View file

@ -1,3 +1,4 @@
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
@ -312,6 +313,61 @@ class TestVertexAIGeminiImageGenerationConfig:
assert result.usage.output_tokens == 17
assert result.usage.total_tokens == 110
def test_transform_image_generation_response_preserves_usage_details_and_traffic_type(self):
mock_response: Final = 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",
}
}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 12,
"candidatesTokenCount": 1300,
"thoughtsTokenCount": 40,
"totalTokenCount": 1352,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}],
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": 10},
{"modality": "IMAGE", "tokenCount": 1290},
],
"trafficType": "ON_DEMAND_PRIORITY",
},
}
mock_response.headers = {}
model_response: Final = ImageResponse()
result: Final = self.config.transform_image_generation_response(
model="gemini-3.1-flash-image",
raw_response=mock_response,
model_response=model_response,
logging_obj=MagicMock(),
request_data={},
optional_params={},
litellm_params={},
encoding=None,
)
assert result.usage.output_tokens == 1340
assert result.usage.input_tokens_details.model_dump() == {"text_tokens": 12, "image_tokens": 0}
assert result.usage.completion_tokens_details == {
"text_tokens": 10,
"image_tokens": 1290,
"reasoning_tokens": 40,
}
assert result.usage.output_tokens_details == result.usage.completion_tokens_details
assert result._hidden_params["provider_specific_fields"] == {"traffic_type": "ON_DEMAND_PRIORITY"}
def test_transform_image_generation_response_multiple_images(self):
"""Test response transformation with multiple images"""
mock_response = MagicMock(spec=httpx.Response)