(fix): token_count calculation during inline data sent during multimodal inputs

This commit is contained in:
Praveen11558 2026-03-10 15:34:22 +05:30 • committed by GitHub
parent 69479ab69b
commit db205ceea5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -159,6 +159,43 @@ def test_inline_pdf_tokens_costed_when_text_tokens_less_than_prompt_tokens():
f"Expected cost={expected_cost} (all 5200 prompt tokens costed), "
f"got cost={cost}. PDF tokens are likely not being costed."
)
def test_inline_pdf_tokens_costed_when_text_both_costed_correctly():
usage = Usage(
completion_tokens=80,
prompt_tokens=6000,
total_tokens=6080,
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=500, cached_tokens=None, text_tokens=50, image_tokens=None
),
)
response = ModelResponse(
usage= usage,
model = "gemini-2.0-flash-001",
)
cost = response_cost_calculator(
reponse_object= response,
model = "gemini-2.0-flash-001",
custom_llm_provider = "vertex_ai",
call_type = "acompletion",
optional_params= {},
)
model_info = litellm.model_cost["gemini-2.0-flash-001"]
input_cost_per_token = model_info["input_cost_per_token"]
input_cost_per_audio_token = model_info["input_cost_per_audio_token"]
output_cost_per_token = model_info["output_cost_per_token"]
expected_input_cost = (
5500 * input_cost_per_token + 500 * input_cost_per_audio_token
)
expected_output_cost = 80 * = output_cost_per_token
expected_cost = expected_input_cost + expected_output_cost
assert cost == pytest.approx(expected_cost), (
f"Expected cost={expected_cost}, "got cost={cost}. "
f"Unaccounted PDF tokens alongside audio are not being costed."
)
def test_reasoning_tokens_gemini_3_1_flash_lite():
"""Test cost calculation for gemini-3.1-flash-lite-preview with reasoning tokens"""