scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

test(cost): cover modality guards and image detection fallbacks

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-15 01:34:17 +00:00
parent 4a8ec7b9d8
commit 27a486e4d3
2 changed files with 107 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -108,6 +108,74 @@ def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None:
assert prompt_cost == pytest.approx(2 * 0.00016)
def test_generic_cost_per_token_prefers_image_per_image_rate() -> None:
model_info: ModelInfo = {
"key": "gemini-embedding-2",
"max_tokens": None,
"max_input_tokens": None,
"max_output_tokens": None,
"input_cost_per_token": 2e-7,
"input_cost_per_image_token": 4.5e-7,
"input_cost_per_image": 0.00012,
"output_cost_per_token": 0.0,
"litellm_provider": "vertex_ai",
"mode": "embedding",
"supported_openai_params": None,
}
usage = Usage(
prompt_tokens=258,
completion_tokens=0,
total_tokens=258,
prompt_tokens_details=PromptTokensDetailsWrapper(
image_tokens=258,
image_count=1,
),
)
prompt_cost, _ = generic_cost_per_token(
model="gemini-embedding-2",
usage=usage,
custom_llm_provider="vertex_ai",
model_info=model_info,
)
assert prompt_cost == pytest.approx(0.00012)
def test_generic_cost_per_token_prefers_video_per_second_rate() -> None:
model_info: ModelInfo = {
"key": "gemini-embedding-2",
"max_tokens": None,
"max_input_tokens": None,
"max_output_tokens": None,
"input_cost_per_token": 2e-7,
"input_cost_per_video_token": 1.2e-5,
"input_cost_per_video_per_second": 0.00079,
"output_cost_per_token": 0.0,
"litellm_provider": "vertex_ai",
"mode": "embedding",
"supported_openai_params": None,
}
usage = Usage(
prompt_tokens=516,
completion_tokens=0,
total_tokens=516,
prompt_tokens_details=PromptTokensDetailsWrapper(
video_tokens=516,
video_length_seconds=2,
),
)
prompt_cost, _ = generic_cost_per_token(
model="gemini-embedding-2",
usage=usage,
custom_llm_provider="vertex_ai",
model_info=model_info,
)
assert prompt_cost == pytest.approx(2 * 0.00079)
def test_missing_cache_read_uses_off_peak_input_rate():
from datetime import datetime, timezone

39
tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py
View file

@ -546,6 +546,45 @@ class TestProcessEmbedContentResponseUsage:
)
assert prompt_cost == pytest.approx(258 * 4.5e-7)
@pytest.mark.parametrize(
"input_value,resolved_files,expected_image_tokens",
[
(GCS_URL, {}, 258),
("gs://my-bucket/clip.mp4", {}, 0),
("gs://my-bucket/unknown.bin", {}, 0),
("files/image-123", {"files/image-123": {"mime_type": "image/jpeg"}}, 258),
("files/missing", {}, 0),
("data:application/octet-stream;base64,abc", {}, 0),
([[IMAGE_DATA_URI]], {}, 258),
([], {}, 0),
],
)
def test_missing_modality_details_classifies_image_inputs(self, input_value, resolved_files, expected_image_tokens):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
"promptTokenCount": 258,
"totalTokenCount": 258,
},
}
result = process_embed_content_response(
input=input_value,
model_response=EmbeddingResponse(),
model=self.MODEL,
response_json=response_json,
resolved_files=resolved_files,
)
assert result.usage.prompt_tokens_details.image_tokens == expected_image_tokens
assert result.usage.prompt_tokens_details.text_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
expected_rate = 4.5e-7 if expected_image_tokens else 2e-7
assert prompt_cost == pytest.approx(258 * expected_rate)
def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self):
response_json = {
"embedding": {"values": [0.1]},