diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 53137b73f8b..5b4a1d30ff0 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1237,6 +1237,14 @@ def _normalize_service_tier(service_tier: object) -> str | None: _BASE_PRICING_SERVICE_TIERS: Final[frozenset[str]] = frozenset({"default", "standard"}) +def _served_image_service_tier(completion_response: object, usage_object: Usage | None) -> str | None: + """Image requests never send a tier upstream, so only the tier the provider reports can price them.""" + return _resolve_billable_service_tier( + requested=None, + served=_served_service_tier(completion_response, usage_object), + ) + + def _resolve_billable_service_tier(requested: object, served: object) -> str | None: """Served tier wins when it names a priced tier or explicitly says base pricing; otherwise the request decides.""" served_lower: Final = served.lower() if isinstance(served, str) else None @@ -2043,6 +2051,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=_served_image_service_tier(completion_response, cost_per_token_usage_object), ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 2eb204d151b..1f6a66f699a 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -2107,6 +2107,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. @@ -2154,28 +2155,22 @@ def calculate_image_response_cost_from_usage( else completion_tokens_details_attr ) - if output_tokens_details is None: - completion_tokens_details = CompletionTokensDetailsWrapper( - text_tokens=0, - image_tokens=completion_tokens, - reasoning_tokens=0, - audio_tokens=0, - ) - else: - text_tokens = _get_token_detail_value(output_tokens_details, "text_tokens") or 0 - image_tokens: Final = _get_token_detail_value(output_tokens_details, "image_tokens") or 0 - audio_tokens: Final = _get_token_detail_value(output_tokens_details, "audio_tokens") or 0 - reasoning_tokens: Final = _get_token_detail_value(output_tokens_details, "reasoning_tokens") or 0 - known_output_tokens: Final = text_tokens + image_tokens + audio_tokens + reasoning_tokens - if completion_tokens > known_output_tokens: - text_tokens += completion_tokens - known_output_tokens - - completion_tokens_details = CompletionTokensDetailsWrapper( - text_tokens=text_tokens, - image_tokens=image_tokens, - reasoning_tokens=reasoning_tokens, - audio_tokens=audio_tokens, - ) + text_tokens: Final = _get_token_detail_value(output_tokens_details, "text_tokens") or 0 + image_tokens: Final = _get_token_detail_value(output_tokens_details, "image_tokens") or 0 + audio_tokens: Final = _get_token_detail_value(output_tokens_details, "audio_tokens") or 0 + reasoning_tokens: Final = _get_token_detail_value(output_tokens_details, "reasoning_tokens") or 0 + has_output_modalities: Final = any( + _get_token_detail_value(output_tokens_details, modality) is not None + for modality in ("text_tokens", "image_tokens", "audio_tokens") + ) + unitemized_tokens: Final = max(completion_tokens - text_tokens - image_tokens - audio_tokens - reasoning_tokens, 0) + unitemized_image_tokens: Final = unitemized_tokens if not has_output_modalities and image_response.data else 0 + completion_tokens_details: Final = CompletionTokensDetailsWrapper( + text_tokens=text_tokens + unitemized_tokens - unitemized_image_tokens, + image_tokens=image_tokens + unitemized_image_tokens, + reasoning_tokens=reasoning_tokens, + audio_tokens=audio_tokens, + ) normalized_usage: Final = Usage( prompt_tokens=prompt_tokens, @@ -2191,6 +2186,7 @@ def calculate_image_response_cost_from_usage( custom_llm_provider=custom_llm_provider, model_info=model_info, vertex_location=vertex_location, + service_tier=service_tier, ) return prompt_cost + completion_cost @@ -2298,6 +2294,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 @@ -2354,6 +2351,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): diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index 3ef403f79ee..3ea8e445a47 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -1,67 +1,131 @@ +from collections.abc import Mapping from typing import Final from pydantic import TypeAdapter +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig 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[str, object]) -> ImageUsage: +def _output_modality_token_details( + candidate_tokens: int, + reasoning_tokens: int, + candidate_tokens_are_inclusive: bool, + modality_tokens: ImageUsageInputTokensDetails, +) -> ImageUsageInputTokensDetails: + non_reasoning_candidate_tokens: Final = max( + candidate_tokens - (reasoning_tokens if candidate_tokens_are_inclusive else 0), + 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 = TypeAdapter(int).validate_python(usage_metadata.get("candidatesTokenCount", 0)) - output_tokens_details: Final = _sum_modality_token_details( + 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, + modality_tokens, + ) + if candidate_token_details + else None + ) - 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} + if output_modality_tokens is not None + else {} + ), + **({"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) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 80244cfc3eb..b5267f9681a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2442,7 +2442,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 @@ -2452,12 +2452,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( diff --git a/litellm/llms/vertex_ai/image_generation/cost_calculator.py b/litellm/llms/vertex_ai/image_generation/cost_calculator.py index df8629aac46..6334c8646b6 100644 --- a/litellm/llms/vertex_ai/image_generation/cost_calculator.py +++ b/litellm/llms/vertex_ai/image_generation/cost_calculator.py @@ -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 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 e1049ef8ece..252054d7330 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -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, 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 @@ -21,8 +22,6 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( ImageObject, ImageResponse, - ImageUsage, - ImageUsageInputTokensDetails, ) if TYPE_CHECKING: @@ -262,27 +261,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, @@ -334,7 +312,23 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) 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 = model_response.hidden_params + 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 = { # rebind-ok: transformation fills the caller-owned response metadata + **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: diff --git a/tests/integration/spend/test_vertex_gemini_image_token_spend.py b/tests/integration/spend/test_vertex_gemini_image_token_spend.py new file mode 100644 index 00000000000..65864fcba61 --- /dev/null +++ b/tests/integration/spend/test_vertex_gemini_image_token_spend.py @@ -0,0 +1,158 @@ +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 +_PRIORITY_IMAGE_RATE: Final = 8e-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 * _PRIORITY_IMAGE_RATE +) + + +def _image_reply(traffic_type: str | None, text_only: bool, breakdown: bool) -> 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": 25 if text_only else _TEXT_OUTPUT_TOKENS + _IMAGE_OUTPUT_TOKENS, + "thoughtsTokenCount": 0 if text_only else _THOUGHTS_TOKENS, + "totalTokenCount": _PROMPT_TOKENS + + (25 if text_only else _TEXT_OUTPUT_TOKENS + _IMAGE_OUTPUT_TOKENS + _THOUGHTS_TOKENS), + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": _PROMPT_TOKENS}], + **( + { + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 25 if text_only else _TEXT_OUTPUT_TOKENS}, + *([] if text_only else [{"modality": "IMAGE", "tokenCount": _IMAGE_OUTPUT_TOKENS}]), + ] + } + if breakdown + else {} + ), + **({"trafficType": traffic_type} if traffic_type is not None else {}), + } + return Reply( + body=json.dumps( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "No image generated"} + if text_only + else {"inlineData": {"mimeType": "image/png", "data": _PNG_B64}}, + ], + }, + "finishReason": "NO_IMAGE" if text_only else "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, + output_cost_per_image_token_priority=_PRIORITY_IMAGE_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", "text_only", "breakdown", "expected_spend"), + [ + pytest.param(None, False, True, _EXPECTED_STANDARD_SPEND, id="standard"), + pytest.param("ON_DEMAND_PRIORITY", False, True, _EXPECTED_PRIORITY_SPEND, id="on-demand-priority"), + pytest.param(None, True, True, _PROMPT_TOKENS * _INPUT_RATE + 25 * _TEXT_OUTPUT_RATE, id="text-only"), + pytest.param( + None, True, False, _PROMPT_TOKENS * _INPUT_RATE + 25 * _TEXT_OUTPUT_RATE, id="text-only-no-breakdown" + ), + pytest.param( + None, + False, + False, + _PROMPT_TOKENS * _INPUT_RATE + + _THOUGHTS_TOKENS * _TEXT_OUTPUT_RATE + + (_TEXT_OUTPUT_TOKENS + _IMAGE_OUTPUT_TOKENS) * _IMAGE_OUTPUT_RATE, + id="image-no-breakdown", + ), + ], +) +def test_vertex_gemini_image_spend_prices_text_thinking_and_served_tier( + gateway: Gateway, + traffic_type: str | None, + text_only: bool, + breakdown: bool, + expected_spend: float, +) -> None: + with wire_server(_image_reply(traffic_type, text_only, breakdown)) 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) diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 0aa76deb9d4..cb9d87622e2 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( apply_off_peak_pricing, apply_provider_cache_read_default, calculate_cache_writing_cost, + calculate_image_response_cost_from_usage, generic_cost_per_token, get_billed_token_rates, get_regional_uplift_multiplier, @@ -27,6 +28,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( from litellm.llms.gemini.image_generation.cost_calculator import ( cost_calculator as gemini_image_generation_cost_calculator, ) +from litellm.llms.gemini.image_usage_transformation import transform_gemini_image_usage from litellm.llms.vertex_ai.image_generation.cost_calculator import ( cost_calculator as vertex_image_generation_cost_calculator, ) @@ -3511,6 +3513,48 @@ def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(_local_model assert text_output_cost + geo.reasoning_cost == pytest.approx(completion_cost) +@pytest.mark.parametrize("has_image", [False, True]) +@pytest.mark.parametrize("thoughts", [0, 40]) +@pytest.mark.parametrize("gemini_usage", [False, True]) +@pytest.mark.parametrize("explicit_breakdown", [False, True]) +def test_image_response_unitemized_output_uses_returned_images( + has_image: bool, thoughts: int, gemini_usage: bool, explicit_breakdown: bool +) -> None: + prices: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "output_cost_per_token": 3e-6, + "output_cost_per_image_token": 7e-5, + } + response: Final = ImageResponse( + data=[ImageObject(b64_json="x")] if has_image else [], + usage=transform_gemini_image_usage( + { + "promptTokenCount": 5, + "candidatesTokenCount": 25, + "thoughtsTokenCount": thoughts, + "totalTokenCount": 30 + thoughts, + **({"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 25}]} if explicit_breakdown else {}), + } + ) + if gemini_usage + else ImageUsage( + input_tokens=5, + input_tokens_details=ImageUsageInputTokensDetails(text_tokens=5, image_tokens=0), + output_tokens=25 + thoughts, + total_tokens=30 + thoughts, + **( + {"output_tokens_details": {"text_tokens": 25, "image_tokens": 0, "reasoning_tokens": thoughts}} + if explicit_breakdown + else {"output_tokens_details": {"reasoning_tokens": thoughts}} + ), + ), + ) + expected_output_rate: Final = 7e-5 if has_image and not explicit_breakdown else 3e-6 + assert calculate_image_response_cost_from_usage( + model="gemini-image-test", image_response=response, custom_llm_provider="gemini", model_info=prices + ) == pytest.approx(5 * 2e-6 + 25 * expected_output_rate + thoughts * 3e-6) + + @pytest.mark.parametrize("details_as_dict", [True, False]) def test_image_response_input_image_tokens_priced_at_image_rate(details_as_dict): """ diff --git a/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index c21dec9b6bb..34826fa0571 100644 --- a/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/unit/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -234,12 +234,12 @@ class TestGeminiImageEditTransformation: assert usage["prompt_tokens"] == 35 assert usage["completion_tokens"] == 1716 assert usage["prompt_tokens_details"]["image_tokens"] == 5 - assert usage["completion_tokens_details"]["image_tokens"] == 1716 + assert usage["completion_tokens_details"] == {} logging_usage = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj=image_response.model_dump() ) - assert logging_usage["completion_tokens_details"]["image_tokens"] == 1716 + assert all(value is None for value in logging_usage["completion_tokens_details"].values()) @pytest.mark.parametrize( ("body", "error", "message"), diff --git a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py index 138aafcbfcd..5d5edc2a49c 100644 --- a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py @@ -294,7 +294,7 @@ def test_gemini_image_generation_preserves_tool_config_side_effect(): } -def test_gemini_image_generation_usage_without_output_details_treats_output_as_image(): +def test_gemini_image_generation_preserves_unspecified_output_modality(): config = GoogleImageGenConfig() raw_response = httpx.Response( status_code=200, @@ -334,8 +334,8 @@ 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 + assert usage["completion_tokens_details"] == {} + assert usage["output_tokens"] == 1716 def test_gemini_image_generation_response_tracks_web_search_requests(): diff --git a/tests/unit/llms/gemini/test_image_usage_transformation.py b/tests/unit/llms/gemini/test_image_usage_transformation.py index f2f1c3d999f..cb64ddb61ac 100644 --- a/tests/unit/llms/gemini/test_image_usage_transformation.py +++ b/tests/unit/llms/gemini/test_image_usage_transformation.py @@ -11,7 +11,6 @@ PROMPT_DETAILS: Final = [{"modality": "TEXT", "tokenCount": 30}, {"modality": "I @pytest.mark.parametrize( ("candidates_details", "expected_output_details"), [ - ({}, {"image_tokens": 1716, "text_tokens": 0}), ( {"candidatesTokensDetails": [{"modality": "IMAGE", "tokenCount": 1120}]}, {"image_tokens": 1120, "text_tokens": 596}, @@ -48,6 +47,72 @@ def test_transform_gemini_image_usage_reports_image_and_chat_style_counts( } +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 + + +@pytest.mark.parametrize("thoughts", [0, 40]) +def test_transform_gemini_image_usage_preserves_unspecified_output_modality(thoughts: int) -> None: + usage: Final = transform_gemini_image_usage( + { + "promptTokenCount": 12, + "candidatesTokenCount": 1300, + "thoughtsTokenCount": thoughts, + "totalTokenCount": 1312 + thoughts, + } + ) + + assert usage.output_tokens == 1300 + thoughts + assert usage.completion_tokens_details == { + "reasoning_tokens": thoughts, + } + assert usage.output_tokens_details == usage.completion_tokens_details + + @pytest.mark.parametrize("count_field", ["promptTokenCount", "totalTokenCount"]) def test_transform_gemini_image_usage_rejects_a_count_that_is_not_a_number(count_field: str): with pytest.raises(ValidationError): 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 65d6e97f8e5..c1c4c5010ed 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 @@ -19,6 +19,7 @@ from litellm.types.utils import ( ImageResponse, ImageUsage, ImageUsageInputTokensDetails, + ModelInfo, PromptTokensDetailsWrapper, Usage, ) @@ -148,3 +149,74 @@ def test_vertex_image_generation_cost_no_web_search_when_absent(monkeypatch): ) assert cost_zero == cost_none + + +@pytest.mark.parametrize("priority_image_rate", [None, 8e-5]) +def test_vertex_image_generation_cost_prices_token_details_by_service_tier(priority_image_rate: float | None) -> 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, + "output_cost_per_image_token_priority": priority_image_rate, + "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 if priority_image_rate is None else priority_image_rate) + ) + assert standard_cost == pytest.approx(expected_standard_cost) + assert priority_cost == pytest.approx(expected_priority_cost) + + +def test_vertex_image_generation_without_usage_keeps_flat_image_price() -> None: + cost: Final = vertex_image_generation_cost_calculator( + model="gemini-scripted-image-tokens", + image_response=ImageResponse(data=[ImageObject(b64_json="image")]), + model_info={"output_cost_per_image": 0.07}, + ) + assert cost == pytest.approx(0.07) diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 10ce24e969c..e1eadb89328 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -15,6 +15,7 @@ from litellm.llms.gemini.image_generation.transformation import GoogleImageGenCo from litellm.llms.vertex_ai.image_generation import ( get_vertex_ai_image_generation_config, ) +from litellm.llms.vertex_ai.image_generation.cost_calculator import cost_calculator from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import ( VertexAIGeminiImageGenerationConfig, ) @@ -410,6 +411,97 @@ class TestVertexAIGeminiImageGenerationConfig: assert result.usage.output_tokens == 17 assert result.usage.total_tokens == 110 + def test_transform_text_only_no_image_response_bills_text_tokens(self) -> None: + result: Final = self.config.transform_image_generation_response( + model="gemini-3.1-flash-image", + raw_response=httpx.Response( + 200, + json={ + "candidates": [ + {"content": {"parts": [{"text": "No image generated"}]}, "finishReason": "NO_IMAGE"} + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 25, + "totalTokenCount": 30, + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 25}], + }, + }, + ), + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.data == [] + assert cost_calculator( + model="gemini-3.1-flash-image", + image_response=result, + model_info={ + "input_cost_per_token": 5e-7, + "output_cost_per_token": 3e-6, + "output_cost_per_image_token": 6e-5, + }, + vertex_location="global", + ) == pytest.approx(5 * 5e-7 + 25 * 3e-6) + + 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) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 07b6111aa67..b991f492bdd 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -2503,6 +2503,36 @@ def test_vertex_gemini_flash_image_generation_regional_costs_uplift_over_global( assert image_cost("us-central1") == pytest.approx(global_cost * 1.10, rel=1e-9) +@pytest.mark.parametrize("requested_tier", ["flex", "priority"]) +def test_vertex_gemini_image_generation_bills_only_the_tier_vertex_served( + _local_model_cost_map: None, requested_tier: str +) -> None: + """Image requests never send service_tier upstream, so a request tier must not change the bill + when Vertex reports no trafficType, while a served ON_DEMAND_PRIORITY still bills priority rates.""" + usage: Final = ImageUsage( + input_tokens=1000, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=1000), + output_tokens=1120, + total_tokens=2120, + ) + + def image_cost(traffic_type: str | None, service_tier: str | None) -> float: + response: Final = ImageResponse(data=[ImageObject(b64_json="img")], usage=usage) + response.hidden_params["provider_specific_fields"] = {"traffic_type": traffic_type} if traffic_type else {} + return completion_cost( + completion_response=response, + model="vertex_ai/gemini-3.1-flash-image", + call_type="image_generation", + optional_params={"service_tier": service_tier} if service_tier else {}, + ) + + standard_cost: Final = image_cost(None, None) + assert standard_cost > 0 + assert image_cost(None, requested_tier) == standard_cost + assert image_cost("ON_DEMAND", requested_tier) == standard_cost + assert image_cost("ON_DEMAND_PRIORITY", requested_tier) > standard_cost + + def test_vertex_uplift_composes_with_above_128k_pricing(monkeypatch): """The regional-endpoint uplift multiplies whatever rate the request priced at, including the above-128k dynamic rates, so a synthetic model carrying both keys