diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 86e1a8fd196..8256b1b4f23 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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 ### diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 29b5f1bd66a..353de4f143d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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): diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index 5f08d8403af..b31f6d134a4 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -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) 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 b5f32d57061..057c1e0e5e4 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 @@ -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( 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 9f3126fd4ff..0a422642bea 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 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: 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..6ae408607f6 --- /dev/null +++ b/tests/integration/spend/test_vertex_gemini_image_token_spend.py @@ -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) diff --git a/tests/unit/llms/gemini/test_image_usage_transformation.py b/tests/unit/llms/gemini/test_image_usage_transformation.py index 54e0668d2d9..7987218118b 100644 --- a/tests/unit/llms/gemini/test_image_usage_transformation.py +++ b/tests/unit/llms/gemini/test_image_usage_transformation.py @@ -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 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..81c70ffe722 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,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) 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 0461c739395..986cff41d30 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 @@ -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)