Include multimodal tokens in tiered pricing threshold

This commit is contained in:
its-amann 2026-06-12 18:29:24 +05:30
parent d79dcc7476
commit ceb8dd8074
2 changed files with 64 additions and 1 deletions

View file

@ -751,7 +751,21 @@ def _get_effective_prompt_tokens_for_tiered_pricing(usage: Usage) -> float:
return prompt_tokens
text_tokens: Final = float(raw_text_tokens)
return text_tokens + cached_tokens + cache_creation
audio_tokens: Final = _loose_token_count(_detail_token(details, "audio_tokens"))
image_tokens: Final = _loose_token_count(_detail_token(details, "image_tokens"))
video_tokens: Final = _loose_token_count(_detail_token(details, "video_tokens"))
detail_total: Final = (
text_tokens
+ audio_tokens
+ image_tokens
+ video_tokens
+ cached_tokens
+ cache_creation
)
modality_cache_tokens: Final = cached_tokens + cache_creation
if modality_cache_tokens > 0 and detail_total > prompt_tokens:
return prompt_tokens
return max(prompt_tokens, detail_total)
def _get_token_base_cost(

View file

@ -3640,6 +3640,55 @@ def test_custom_pricing_does_not_double_count_cached_tokens_when_text_tokens_mis
litellm.model_cost.pop(f"anthropic/{model_id}", None)
def test_custom_pricing_uses_tiered_rate_when_multimodal_details_cross_threshold():
"""Image tokens counted beside text_tokens cross the above-200k tier."""
model_id = "tiered-custom-pricing-multimodal-details-test"
model_info = {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 0.000005,
"output_cost_per_token": 0.000025,
"input_cost_per_token_above_200k_tokens": 0.00001,
"output_cost_per_token_above_200k_tokens": 0.0000375,
}
text_tokens = 150_000
image_tokens = 60_001
litellm.register_model(model_cost={model_id: model_info})
try:
response = ModelResponse(
id="test-id",
created=1234567890,
model=f"openai/{model_id}",
object="chat.completion",
choices=[],
usage=Usage(
prompt_tokens=text_tokens + image_tokens,
completion_tokens=0,
total_tokens=text_tokens + image_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=text_tokens, image_tokens=image_tokens),
),
)
cost = completion_cost(
completion_response=response,
model=f"openai/{model_id}",
custom_llm_provider="openai",
custom_pricing=True,
router_model_id=model_id,
)
expected_cost = (text_tokens + image_tokens) * model_info["input_cost_per_token_above_200k_tokens"]
base_cost = (text_tokens + image_tokens) * model_info["input_cost_per_token"]
assert cost == pytest.approx(expected_cost)
assert cost != pytest.approx(base_cost)
finally:
litellm.model_cost.pop(model_id, None)
litellm.model_cost.pop(f"openai/{model_id}", None)
# ---------------------------------------------------------------------------
# Bug 2 — db_spend_update_writer cache token extraction helpers.
# ---------------------------------------------------------------------------