From 6e3540856cd1b3527f58e6bd1649ef2a6cb321e4 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 19:53:16 -0700 Subject: [PATCH] fix(vertex): preserve Gemini Embedding 2 usageMetadata for cost tracking (#31354) * fix(vertex): preserve Gemini Embedding 2 usageMetadata for cost tracking * style(vertex): apply ruff format to batch_embed_content_transformation * fix(vertex): bill files/ image refs in Gemini embedContent at per-image rate Resolved files/... references whose mime type is an image were not detected by _is_image_element, so image_count stayed 0 and generic_cost_per_token fell back to the text token rate instead of input_cost_per_image. Thread the resolved_files mapping into the usage builder so resolved image references are counted and billed per image. Also modernize the _flatten_input return annotation to satisfy the ruff UP006 strict gate. * fix(vertex): bill Gemini embedding audio per-second and stop video+audio double-billing Audio-only embedContent responses set audio_tokens, but generic_cost_per_token only charges audio via input_cost_per_audio_token. gemini-embedding-2 prices audio via input_cost_per_audio_per_second, so spend stayed at $0. Plumb a new audio_length_seconds field through PromptTokensDetailsWrapper, parse it in _parse_prompt_tokens_details, and bill it from _calculate_input_cost. The vertex embedding transformation derives audio_length_seconds from audio_tokens using the documented 32 tokens/sec Gemini rate. The 1-token text floor that protects video billing only fired when no other modality was billable, but audio presence flipped that flag, leaving text_tokens at zero for video+audio responses. generic_cost_per_token then rewrote text_tokens to prompt_tokens minus audio_tokens (the video token count), charging video tokens as text on top of the per-second video cost. The rewrite trigger is text_tokens == 0 and image_count == 0; align the floor with that trigger and ignore audio_tokens. --------- Co-authored-by: Cursor Agent --- .../litellm_core_utils/llm_cost_calc/utils.py | 18 ++ .../batch_embed_content_handler.py | 2 + .../batch_embed_content_transformation.py | 151 +++++++++++- litellm/types/utils.py | 5 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 1 + ...test_batch_embed_content_transformation.py | 219 +++++++++++++++++- 6 files changed, 383 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 19b711f6748..d2c97a5d6e7 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -485,6 +485,7 @@ class PromptTokensDetailsResult(TypedDict): character_count: int image_count: int video_length_seconds: float + audio_length_seconds: float def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -535,6 +536,13 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + audio_length_seconds = ( + cast( + Optional[float], + getattr(usage.prompt_tokens_details, "audio_length_seconds", 0), + ) + or 0.0 + ) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -546,6 +554,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: character_count=character_count, image_count=image_count, video_length_seconds=float(video_length_seconds), + audio_length_seconds=float(audio_length_seconds), ) @@ -667,6 +676,14 @@ def _calculate_input_cost( prompt_tokens_details["video_length_seconds"], ) + ### AUDIO LENGTH COST + if prompt_tokens_details["audio_length_seconds"]: + prompt_cost += calculate_cost_component( + model_info, + "input_cost_per_audio_per_second", + prompt_tokens_details["audio_length_seconds"], + ) + return prompt_cost @@ -743,6 +760,7 @@ def generic_cost_per_token( character_count=0, image_count=0, video_length_seconds=0.0, + audio_length_seconds=0.0, ) if usage.prompt_tokens_details: prompt_tokens_details = _parse_prompt_tokens_details(usage) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 165dac24903..53f5fb464df 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -274,6 +274,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore @@ -377,6 +378,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions = VertexAIBatchEmbeddingsResponseObject(**_json_response) # type: ignore diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index ba6e6f0c056..27ca1bd92a4 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,10 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from typing import Dict, List, Optional, Tuple +from collections.abc import Mapping +from typing import Dict, List, Optional, Sequence, Tuple + +from pydantic import TypeAdapter, ValidationError from litellm.types.llms.vertex_ai import ( BlobType, @@ -13,10 +16,17 @@ from litellm.types.llms.vertex_ai import ( FileDataType, GeminiEmbeddingInput, PartType, + PromptTokensDetails, + UsageMetadata, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import ( + Embedding, + EmbeddingResponse, + PromptTokensDetailsWrapper, + Usage, +) from litellm.utils import get_formatted_prompt, token_counter SUPPORTED_EMBEDDING_MIME_TYPES = { @@ -294,11 +304,133 @@ def transform_openai_input_gemini_embed_content( return request_body +_IMAGE_MIME_TYPES = frozenset({"image/png", "image/jpeg"}) +_VIDEO_TOKENS_PER_SECOND = 258.0 +_AUDIO_TOKENS_PER_SECOND = 32.0 +_usage_metadata_adapter = TypeAdapter(UsageMetadata) + + +def _parse_usage_metadata(raw_usage_metadata: object) -> Optional[UsageMetadata]: + if not isinstance(raw_usage_metadata, dict): + return None + try: + return _usage_metadata_adapter.validate_python(raw_usage_metadata) + except ValidationError: + 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 = 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 + ) + + +def _fallback_usage(input: GeminiEmbeddingInput, model: str) -> Usage: + if _is_multimodal_input(input): + return Usage(prompt_tokens=0, total_tokens=0) + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) + return Usage(prompt_tokens=prompt_tokens, total_tokens=prompt_tokens) + + +def _usage_from_embed_content_response( + input: GeminiEmbeddingInput, + model: str, + raw_usage_metadata: object, + resolved_files: Mapping[str, Mapping[str, str]], +) -> Usage: + usage_metadata = _parse_usage_metadata(raw_usage_metadata) + if usage_metadata is None: + return _fallback_usage(input, model) + + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + total_tokens = usage_metadata.get("totalTokenCount") or prompt_tokens + + details: Sequence[PromptTokensDetails] = ( + usage_metadata.get("promptTokensDetails") or () + ) + text_tokens = _tokens_for_modality(details, "TEXT") + audio_tokens = _tokens_for_modality(details, "AUDIO") + video_tokens = _tokens_for_modality(details, "VIDEO") + image_count = _count_input_images(input, resolved_files) + + video_length_seconds = ( + video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 + ) + audio_length_seconds = ( + 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 = ( + video_length_seconds > 0 and text_tokens == 0 and image_count == 0 + ) + resolved_text_tokens = 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, + audio_tokens=audio_tokens, + image_count=image_count, + video_length_seconds=video_length_seconds, + audio_length_seconds=audio_length_seconds, + ), + ) + + def process_embed_content_response( input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, response_json: dict, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). @@ -308,6 +440,8 @@ 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 @@ -327,14 +461,11 @@ def process_embed_content_response( model_response.data = [openai_embedding] model_response.model = model - - if _is_multimodal_input(input): - prompt_tokens = 0 - else: - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) - model_response.usage = Usage( - prompt_tokens=prompt_tokens, total_tokens=prompt_tokens + model_response.usage = _usage_from_embed_content_response( + input=input, + model=model, + raw_usage_metadata=response_json.get("usageMetadata"), + resolved_files=resolved_files or {}, ) return model_response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 24d6e84fba7..c3f99b1d18b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1539,6 +1539,9 @@ class PromptTokensDetailsWrapper( video_length_seconds: Optional[float] = None """Length of videos sent to the model. Used for Vertex AI multimodal embeddings.""" + audio_length_seconds: Optional[float] = None + """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + cache_creation_tokens: Optional[int] = None """Number of cache creation tokens sent to the model. Used for Anthropic prompt caching.""" @@ -1553,6 +1556,8 @@ class PromptTokensDetailsWrapper( del self.image_count if self.video_length_seconds is None: del self.video_length_seconds + if self.audio_length_seconds is None: + del self.audio_length_seconds if self.web_search_requests is None: del self.web_search_requests if self.cache_creation_tokens is None: diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 7f3d5a959a1..a5d1934237f 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -951,6 +951,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "character_count": 0, "image_count": 0, "video_length_seconds": 0.0, + "audio_length_seconds": 0.0, } model_info: ModelInfo = {} diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index bb4e6c67e9e..86b3f0976ab 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -10,9 +10,11 @@ Covers: import pytest +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, _is_multimodal_input, + process_embed_content_response, process_response, transform_openai_input_gemini_content, transform_openai_input_gemini_embed_content, @@ -72,7 +74,9 @@ class TestBuildPartForInput: assert part["file_data"]["file_uri"] == GCS_URL def test_file_reference_resolved(self): - resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} + resolved = { + "files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"} + } part = _build_part_for_input("files/abc", resolved_files=resolved) assert part["file_data"] is not None assert part["file_data"]["mime_type"] == "image/jpeg" @@ -94,7 +98,9 @@ class TestTransformOpenaiInputGeminiContent: def test_multiple_texts(self): result = transform_openai_input_gemini_content( - input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={} + input=["hello", "world"], + model="gemini-embedding-2-preview", + optional_params={}, ) assert len(result["requests"]) == 2 assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" @@ -109,7 +115,10 @@ class TestTransformOpenaiInputGeminiContent: ) assert len(result["requests"]) == 2 # First request is text - assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" + assert ( + result["requests"][0]["content"]["parts"][0]["text"] + == "The food was delicious" + ) # Second request is image assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None @@ -288,3 +297,207 @@ class TestProcessResponse: model="gemini-embedding-2-preview", optional_params={}, ) + + +class TestProcessEmbedContentResponseUsage: + """Gemini Embedding 2 embedContent usageMetadata must drive spend. + + Regression for multimodal calls recording prompt_tokens=0 / spend=$0. + """ + + MODEL = "gemini-embedding-2" + + def test_multimodal_image_preserves_usage_metadata(self): + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 258 + assert result.usage.total_tokens == 258 + assert result.usage.prompt_tokens_details.image_count == 1 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_text_modality_detail_populated(self): + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + }, + } + result = process_embed_content_response( + input="a short caption", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 12 + assert result.usage.prompt_tokens_details.text_tokens == 12 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost > 0 + + def test_video_modality_derives_seconds_and_text_floor(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 516, + "totalTokenCount": 516, + "promptTokensDetails": [{"modality": "VIDEO", "tokenCount": 516}], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + model=self.MODEL, + 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 + + def test_missing_usage_metadata_does_not_estimate_from_base64(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input=[IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens == 0 + assert result.usage.total_tokens == 0 + + def test_missing_usage_metadata_text_falls_back_to_token_counter(self): + response_json = {"embedding": {"values": [0.1, 0.2]}} + result = process_embed_content_response( + input="hello world this is plain text", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + 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.""" + response_json = { + "embedding": {"values": [0.1, 0.2, 0.3]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + "promptTokensDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + } + result = process_embed_content_response( + input=["files/img123"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/img123": { + "mime_type": "image/png", + "uri": "https://example.com/img123", + } + }, + ) + assert result.usage.prompt_tokens_details.image_count == 1 + 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", + ) + assert prompt_cost == pytest.approx(0.00012) + + 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.""" + response_json = { + "embedding": {"values": [0.1, 0.2]}, + "usageMetadata": { + "promptTokenCount": 64, + "totalTokenCount": 64, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}], + }, + } + result = process_embed_content_response( + input=["files/clip1"], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files={ + "files/clip1": { + "mime_type": "audio/mpeg", + "uri": "https://example.com/clip1", + } + }, + ) + 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 + ) + + 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) + + def test_video_plus_audio_does_not_double_bill_text(self): + """Video+audio responses must not get video tokens reassigned to text.""" + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 580, + "totalTokenCount": 580, + "promptTokensDetails": [ + {"modality": "VIDEO", "tokenCount": 516}, + {"modality": "AUDIO", "tokenCount": 64}, + ], + }, + } + result = process_embed_content_response( + input=["gs://bucket/clip.mp4"], + model_response=EmbeddingResponse(), + 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 + ) + + 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)