mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
refactor(cost): dedupe image token usage cost helper
- extract shared calculate_image_response_cost_from_usage() helper\n- reuse helper in vertex and gemini image generation cost calculators\n- preserve provider-specific fallback to output_cost_per_image
This commit is contained in:
Emerson Gomes
parent
702d5e88b8
commit
0e014253d7
3 changed files with 76 additions and 130 deletions
|
|
@ -8,9 +8,11 @@ from litellm._logging import verbose_logger
|
|||
from litellm.types.utils import (
|
||||
CacheCreationTokenDetails,
|
||||
CallTypes,
|
||||
CompletionTokensDetailsWrapper,
|
||||
ImageResponse,
|
||||
ModelInfo,
|
||||
PassthroughCallTypes,
|
||||
PromptTokensDetailsWrapper,
|
||||
ServiceTier,
|
||||
Usage,
|
||||
)
|
||||
|
|
@ -767,6 +769,64 @@ def generic_cost_per_token( # noqa: PLR0915
|
|||
return prompt_cost, completion_cost
|
||||
|
||||
|
||||
def calculate_image_response_cost_from_usage(
|
||||
model: str,
|
||||
image_response: ImageResponse,
|
||||
custom_llm_provider: str,
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
Calculate image generation cost from usage metadata when available.
|
||||
|
||||
Returns:
|
||||
Optional[float]: total cost from token usage, or None when usage metadata
|
||||
is missing/incomplete and caller should fall back to flat per-image pricing.
|
||||
"""
|
||||
usage = image_response.usage
|
||||
if usage is None:
|
||||
return None
|
||||
|
||||
prompt_tokens = usage.input_tokens
|
||||
completion_tokens = usage.output_tokens
|
||||
total_tokens = usage.total_tokens
|
||||
|
||||
if prompt_tokens is None or completion_tokens is None or total_tokens is None:
|
||||
return None
|
||||
|
||||
# ImageResponse may carry a default zeroed usage object even when provider
|
||||
# usage metadata is absent. Treat this as missing usage and fall back.
|
||||
if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
|
||||
return None
|
||||
|
||||
input_tokens_details = getattr(usage, "input_tokens_details", None)
|
||||
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
|
||||
if input_tokens_details is not None:
|
||||
prompt_tokens_details = PromptTokensDetailsWrapper(
|
||||
text_tokens=getattr(input_tokens_details, "text_tokens", None),
|
||||
image_tokens=getattr(input_tokens_details, "image_tokens", None),
|
||||
cached_tokens=0,
|
||||
)
|
||||
|
||||
normalized_usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=completion_tokens,
|
||||
reasoning_tokens=0,
|
||||
audio_tokens=0,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model,
|
||||
usage=normalized_usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
|
||||
|
||||
class CostCalculatorUtils:
|
||||
@staticmethod
|
||||
def _call_type_has_image_response(call_type: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -2,71 +2,13 @@
|
|||
Google AI Image Generation Cost Calculator
|
||||
"""
|
||||
|
||||
from typing import Any, Optional
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
ImageResponse,
|
||||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
calculate_image_response_cost_from_usage,
|
||||
)
|
||||
|
||||
|
||||
def _calculate_token_based_cost(model: str, image_response: ImageResponse) -> Optional[float]:
|
||||
"""
|
||||
Calculate token-based image generation cost when usage metadata is available.
|
||||
|
||||
Falls back to None when usage metadata is missing/incomplete.
|
||||
"""
|
||||
usage = image_response.usage
|
||||
if usage is None:
|
||||
return None
|
||||
|
||||
prompt_tokens = usage.input_tokens
|
||||
completion_tokens = usage.output_tokens
|
||||
total_tokens = usage.total_tokens
|
||||
|
||||
if (
|
||||
prompt_tokens is None
|
||||
or completion_tokens is None
|
||||
or total_tokens is None
|
||||
):
|
||||
return None
|
||||
# ImageResponse may carry a default zeroed usage object even when provider
|
||||
# usage metadata is absent. Treat this as missing usage and fall back.
|
||||
if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
|
||||
return None
|
||||
|
||||
input_tokens_details = getattr(usage, "input_tokens_details", None)
|
||||
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
|
||||
if input_tokens_details is not None:
|
||||
prompt_tokens_details = PromptTokensDetailsWrapper(
|
||||
text_tokens=getattr(input_tokens_details, "text_tokens", None),
|
||||
image_tokens=getattr(input_tokens_details, "image_tokens", None),
|
||||
cached_tokens=0,
|
||||
)
|
||||
|
||||
normalized_usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=completion_tokens,
|
||||
reasoning_tokens=0,
|
||||
audio_tokens=0,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model,
|
||||
usage=normalized_usage,
|
||||
custom_llm_provider="gemini",
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
|
||||
def cost_calculator(
|
||||
|
|
@ -74,7 +16,7 @@ def cost_calculator(
|
|||
image_response: Any,
|
||||
) -> float:
|
||||
"""
|
||||
Vertex AI Image Generation Cost Calculator
|
||||
Google AI Image Generation Cost Calculator
|
||||
"""
|
||||
_model_info = litellm.get_model_info(
|
||||
model=model,
|
||||
|
|
@ -82,8 +24,10 @@ def cost_calculator(
|
|||
)
|
||||
|
||||
if isinstance(image_response, ImageResponse):
|
||||
token_based_cost = _calculate_token_based_cost(
|
||||
model=model, image_response=image_response
|
||||
token_based_cost = calculate_image_response_cost_from_usage(
|
||||
model=model,
|
||||
image_response=image_response,
|
||||
custom_llm_provider="gemini",
|
||||
)
|
||||
if token_based_cost is not None:
|
||||
return token_based_cost
|
||||
|
|
|
|||
|
|
@ -2,71 +2,11 @@
|
|||
Vertex AI Image Generation Cost Calculator
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
ImageResponse,
|
||||
PromptTokensDetailsWrapper,
|
||||
Usage,
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
calculate_image_response_cost_from_usage,
|
||||
)
|
||||
|
||||
|
||||
def _calculate_token_based_cost(model: str, image_response: ImageResponse) -> Optional[float]:
|
||||
"""
|
||||
Calculate token-based image generation cost when usage metadata is available.
|
||||
|
||||
Falls back to None when usage metadata is missing/incomplete.
|
||||
"""
|
||||
usage = image_response.usage
|
||||
if usage is None:
|
||||
return None
|
||||
|
||||
prompt_tokens = usage.input_tokens
|
||||
completion_tokens = usage.output_tokens
|
||||
total_tokens = usage.total_tokens
|
||||
|
||||
if (
|
||||
prompt_tokens is None
|
||||
or completion_tokens is None
|
||||
or total_tokens is None
|
||||
):
|
||||
return None
|
||||
# ImageResponse may carry a default zeroed usage object even when provider
|
||||
# usage metadata is absent. Treat this as missing usage and fall back.
|
||||
if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
|
||||
return None
|
||||
|
||||
input_tokens_details = getattr(usage, "input_tokens_details", None)
|
||||
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
|
||||
if input_tokens_details is not None:
|
||||
prompt_tokens_details = PromptTokensDetailsWrapper(
|
||||
text_tokens=getattr(input_tokens_details, "text_tokens", None),
|
||||
image_tokens=getattr(input_tokens_details, "image_tokens", None),
|
||||
cached_tokens=0,
|
||||
)
|
||||
|
||||
normalized_usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=completion_tokens,
|
||||
reasoning_tokens=0,
|
||||
audio_tokens=0,
|
||||
),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model,
|
||||
usage=normalized_usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
|
||||
def cost_calculator(
|
||||
|
|
@ -81,8 +21,10 @@ def cost_calculator(
|
|||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
token_based_cost = _calculate_token_based_cost(
|
||||
model=model, image_response=image_response
|
||||
token_based_cost = calculate_image_response_cost_from_usage(
|
||||
model=model,
|
||||
image_response=image_response,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
if token_based_cost is not None:
|
||||
return token_based_cost
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue