fix(cost): bill gemini-embedding-2 per token and stop double charging audio

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-15 00:41:46 +00:00
parent b67137b67f
commit d4f2119b03
6 changed files with 75 additions and 104 deletions

View file

@ -956,12 +956,17 @@ def _calculate_input_cost(
)
### AUDIO COST
if prompt_tokens_details["audio_tokens"]:
if prompt_tokens_details["audio_tokens"] and not (
prompt_tokens_details["audio_length_seconds"]
and model_info.get("input_cost_per_audio_per_second") is not None
):
audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
### IMAGE TOKEN COST
if prompt_tokens_details["image_tokens"]:
if prompt_tokens_details["image_tokens"] and not (
prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
):
# For image token costs:
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
image_token_cost_key = "input_cost_per_image_token"
@ -970,7 +975,10 @@ def _calculate_input_cost(
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
### VIDEO TOKEN COST
if prompt_tokens_details["video_tokens"]:
if prompt_tokens_details["video_tokens"] and not (
prompt_tokens_details["video_length_seconds"]
and model_info.get("input_cost_per_video_per_second") is not None
):
video_token_cost_key = "input_cost_per_video_token"
if model_info.get(video_token_cost_key) is None:
video_token_cost_key = "input_cost_per_token"

View file

@ -297,9 +297,6 @@ def transform_openai_input_gemini_embed_content(
return request_body
_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"})
_VIDEO_TOKENS_PER_SECOND: Final = 258.0
_AUDIO_TOKENS_PER_SECOND: Final = 32.0
_usage_metadata_adapter: Final = TypeAdapter(UsageMetadata)
@ -312,40 +309,6 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None:
return None
def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]:
if isinstance(input, str):
return (input,)
return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element]))
def _is_image_element(
element: str,
resolved_files: Mapping[str, Mapping[str, str]],
) -> bool:
if element.startswith("data:") and ";base64," in element:
try:
mime_type, _ = _parse_data_url(element)
except ValueError:
return False
return mime_type in _IMAGE_MIME_TYPES
if _is_gcs_url(element):
try:
return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES
except ValueError:
return False
if _is_file_reference(element):
file_info: Final = resolved_files.get(element)
return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES
return False
def _count_input_images(
input: GeminiEmbeddingInput,
resolved_files: Mapping[str, Mapping[str, str]],
) -> int:
return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files))
def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
return sum(detail["tokenCount"] for detail in details if detail["modality"] == modality)
@ -362,7 +325,6 @@ def _usage_from_embed_content_response(
input: GeminiEmbeddingInput,
model: str,
raw_usage_metadata: object,
resolved_files: Mapping[str, Mapping[str, str]],
) -> Usage:
usage_metadata: Final = _parse_usage_metadata(raw_usage_metadata)
if usage_metadata is None:
@ -374,28 +336,17 @@ def _usage_from_embed_content_response(
details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or ()
text_tokens: Final = _tokens_for_modality(details, "TEXT")
audio_tokens: Final = _tokens_for_modality(details, "AUDIO")
image_tokens: Final = _tokens_for_modality(details, "IMAGE")
video_tokens: Final = _tokens_for_modality(details, "VIDEO")
image_count: Final = _count_input_images(input, resolved_files)
video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
# generic_cost_per_token rewrites text_tokens to the full prompt minus
# other modalities when both text_tokens and image_count are zero. For
# video, that misallocates video tokens to text; a 1-token floor sidesteps
# the rewrite and keeps billing on input_cost_per_video_per_second.
needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0
resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens
return Usage(
prompt_tokens=prompt_tokens,
total_tokens=total_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=resolved_text_tokens,
text_tokens=text_tokens,
audio_tokens=audio_tokens,
image_count=image_count,
video_length_seconds=video_length_seconds,
audio_length_seconds=audio_length_seconds,
image_tokens=image_tokens,
video_tokens=video_tokens,
),
)
@ -415,8 +366,6 @@ def process_embed_content_response(
model_response: EmbeddingResponse to populate
model: Model name
response_json: Raw JSON response from embedContent endpoint
resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
used to bill resolved image references at the per-image rate
Returns:
EmbeddingResponse with single embedding
@ -438,7 +387,6 @@ def process_embed_content_response(
input=input,
model=model,
raw_usage_metadata=response_json.get("usageMetadata"),
resolved_files=resolved_files or {},
)
return model_response

View file

@ -25615,13 +25615,11 @@
"uses_embed_content": true
},
"gemini-embedding-2": {
"input_cost_per_audio_per_second": 0.00016,
"input_cost_per_audio_token": 6.5e-06,
"input_cost_per_image": 0.00012,
"input_cost_per_image_token": 4.5e-07,
"input_cost_per_token": 2e-07,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_video_per_second": 0.00079,
"input_cost_per_video_token": 1.2e-05,
"litellm_provider": "vertex_ai-embedding-models",
"max_input_tokens": 8192,
"max_tokens": 8192,
@ -25648,13 +25646,11 @@
"uses_embed_content": true
},
"vertex_ai/gemini-embedding-2": {
"input_cost_per_audio_per_second": 0.00016,
"input_cost_per_audio_token": 6.5e-06,
"input_cost_per_image": 0.00012,
"input_cost_per_image_token": 4.5e-07,
"input_cost_per_token": 2e-07,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_video_per_second": 0.00079,
"input_cost_per_video_token": 1.2e-05,
"litellm_provider": "vertex_ai",
"max_input_tokens": 8192,
"max_tokens": 8192,
@ -25709,10 +25705,11 @@
"tpm": 10000000
},
"gemini/gemini-embedding-2": {
"input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012,
"input_cost_per_audio_token": 6.5e-06,
"input_cost_per_image_token": 4.5e-07,
"input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.00079,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_video_token": 1.2e-05,
"litellm_provider": "gemini",
"max_input_tokens": 8192,
"max_tokens": 8192,

View file

@ -25615,13 +25615,11 @@
"uses_embed_content": true
},
"gemini-embedding-2": {
"input_cost_per_audio_per_second": 0.00016,
"input_cost_per_audio_token": 6.5e-06,
"input_cost_per_image": 0.00012,
"input_cost_per_image_token": 4.5e-07,
"input_cost_per_token": 2e-07,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_video_per_second": 0.00079,
"input_cost_per_video_token": 1.2e-05,
"litellm_provider": "vertex_ai-embedding-models",
"max_input_tokens": 8192,
"max_tokens": 8192,
@ -25648,13 +25646,11 @@
"uses_embed_content": true
},
"vertex_ai/gemini-embedding-2": {
"input_cost_per_audio_per_second": 0.00016,
"input_cost_per_audio_token": 6.5e-06,
"input_cost_per_image": 0.00012,
"input_cost_per_image_token": 4.5e-07,
"input_cost_per_token": 2e-07,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_video_per_second": 0.00079,
"input_cost_per_video_token": 1.2e-05,
"litellm_provider": "vertex_ai",
"max_input_tokens": 8192,
"max_tokens": 8192,
@ -25709,10 +25705,11 @@
"tpm": 10000000
},
"gemini/gemini-embedding-2": {
"input_cost_per_audio_per_second": 0.00016,
"input_cost_per_image": 0.00012,
"input_cost_per_audio_token": 6.5e-06,
"input_cost_per_image_token": 4.5e-07,
"input_cost_per_token": 2e-07,
"input_cost_per_video_per_second": 0.00079,
"input_cost_per_token_batches": 1e-07,
"input_cost_per_video_token": 1.2e-05,
"litellm_provider": "gemini",
"max_input_tokens": 8192,
"max_tokens": 8192,

View file

@ -74,6 +74,39 @@ def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, s
assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4])
def test_generic_cost_per_token_prefers_audio_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_audio_token": 6.5e-6,
"input_cost_per_audio_per_second": 0.00016,
"output_cost_per_token": 0.0,
"litellm_provider": "vertex_ai",
"mode": "embedding",
}
usage = Usage(
prompt_tokens=64,
completion_tokens=0,
total_tokens=64,
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=64,
audio_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.00016)
def test_missing_cache_read_uses_off_peak_input_rate():
from datetime import datetime, timezone

View file

@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation
from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject
from litellm.types.utils import EmbeddingResponse
IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
GCS_URL = "gs://my-bucket/image.png"
@ -324,7 +323,7 @@ class TestProcessEmbedContentResponseUsage:
)
assert result.usage.prompt_tokens == 258
assert result.usage.total_tokens == 258
assert result.usage.prompt_tokens_details.image_count == 1
assert result.usage.prompt_tokens_details.image_tokens == 258
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
@ -358,7 +357,7 @@ class TestProcessEmbedContentResponseUsage:
)
assert prompt_cost > 0
def test_video_modality_derives_seconds_and_text_floor(self):
def test_video_modality_preserves_token_count(self):
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
@ -374,10 +373,8 @@ class TestProcessEmbedContentResponseUsage:
response_json=response_json,
)
assert result.usage.prompt_tokens == 516
assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(
2.0
)
assert result.usage.prompt_tokens_details.text_tokens == 1
assert result.usage.prompt_tokens_details.video_tokens == 516
assert result.usage.prompt_tokens_details.text_tokens == 0
def test_missing_usage_metadata_does_not_estimate_from_base64(self):
response_json = {"embedding": {"values": [0.1, 0.2]}}
@ -400,8 +397,7 @@ class TestProcessEmbedContentResponseUsage:
)
assert result.usage.prompt_tokens > 0
def test_file_reference_image_billed_per_image_not_text(self):
"""files/... image refs must bill per-image, not at the text token rate."""
def test_file_reference_image_billed_per_image_token_rate(self):
response_json = {
"embedding": {"values": [0.1, 0.2, 0.3]},
"usageMetadata": {
@ -422,7 +418,7 @@ class TestProcessEmbedContentResponseUsage:
}
},
)
assert result.usage.prompt_tokens_details.image_count == 1
assert result.usage.prompt_tokens_details.image_tokens == 258
assert result.usage.prompt_tokens_details.text_tokens == 0
prompt_cost, _ = generic_cost_per_token(
@ -430,10 +426,10 @@ class TestProcessEmbedContentResponseUsage:
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(0.00012)
assert prompt_cost == pytest.approx(258 * 4.5e-7)
def test_file_reference_non_image_not_counted_as_image(self):
"""A files/... ref resolving to a non-image mime must not be image-counted."""
"""A files/... ref resolving to a non-image mime keeps audio token billing."""
response_json = {
"embedding": {"values": [0.1, 0.2]},
"usageMetadata": {
@ -454,21 +450,18 @@ class TestProcessEmbedContentResponseUsage:
}
},
)
assert result.usage.prompt_tokens_details.image_count == 0
assert result.usage.prompt_tokens_details.audio_tokens == 64
assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(
2.0
)
assert result.usage.prompt_tokens_details.image_tokens == 0
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
assert prompt_cost == pytest.approx(2.0 * 0.00016)
assert prompt_cost == pytest.approx(64 * 6.5e-6)
def test_video_plus_audio_does_not_double_bill_text(self):
"""Video+audio responses must not get video tokens reassigned to text."""
"""Video and audio responses are billed from their respective token counts."""
response_json = {
"embedding": {"values": [0.1]},
"usageMetadata": {
@ -486,18 +479,13 @@ class TestProcessEmbedContentResponseUsage:
model=self.MODEL,
response_json=response_json,
)
assert result.usage.prompt_tokens_details.text_tokens == 1
assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx(
2.0
)
assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx(
2.0
)
assert result.usage.prompt_tokens_details.text_tokens == 0
assert result.usage.prompt_tokens_details.video_tokens == 516
assert result.usage.prompt_tokens_details.audio_tokens == 64
prompt_cost, _ = generic_cost_per_token(
model=self.MODEL,
usage=result.usage,
custom_llm_provider="vertex_ai",
)
# 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4
assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016)
assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6)