Break Azure AI image cost import cycle

This commit is contained in:
Emerson Gomes 2026-05-02 15:32:25 -05:00
parent 5668e26190
commit e4172c2e93
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17

View file

@ -1,9 +1,6 @@
from typing import Any
import litellm
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_image_response_cost_from_usage,
)
from litellm.types.utils import ImageResponse
@ -23,10 +20,9 @@ def cost_calculator(
f"image_response must be of type ImageResponse got type={type(image_response)}"
)
usage_based_cost = calculate_image_response_cost_from_usage(
usage_based_cost = _calculate_cost_from_usage(
model=model,
image_response=image_response,
custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
)
if usage_based_cost is not None:
return usage_based_cost
@ -38,3 +34,34 @@ def cost_calculator(
output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
num_images = len(image_response.data) if image_response.data else 0
return output_cost_per_image * num_images
def _calculate_cost_from_usage(
model: str,
image_response: ImageResponse,
) -> float | None:
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
if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
return None
model_info = litellm.get_model_info(
model=model,
custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
)
input_cost_per_token = model_info.get("input_cost_per_token") or 0.0
output_cost_per_image_token = (
model_info.get("output_cost_per_image_token")
or model_info.get("output_cost_per_token")
or 0.0
)
return (float(prompt_tokens) * input_cost_per_token) + (
float(completion_tokens) * output_cost_per_image_token
)