mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(vertex_ai): bill gemini image text, thinking and priority tokens at their own rates (#51017)
* fix(vertex_ai): bill gemini image text, thinking and priority tokens at their own rates Co-authored-by: kerry <kerry@berri.ai> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(images): bill unitemized output as text when no images return Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost_calculator): price image calls only by the tier the provider served --------- Co-authored-by: kenneth <kenneth@berri.ai> Co-authored-by: kerry <kerry@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
bb89296adf
commit
461b56b14b
14 changed files with 628 additions and 95 deletions
|
|
@ -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 ###
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
158
tests/integration/spend/test_vertex_gemini_image_token_spend.py
Normal file
158
tests/integration/spend/test_vertex_gemini_image_token_spend.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue