From c6616a3e2f1fa95a38ac5250bc4d0e1fbdde0173 Mon Sep 17 00:00:00 2001 From: omarboukherys Date: Mon, 2 Nov 2026 15:55:54 +0000 Subject: [PATCH] fix(gemini): bill batch embeddings from usageMetadata so media inputs are not free --- .../batch_embed_content_handler.py | 2 + .../batch_embed_content_transformation.py | 10 +++- .../vertex_passthrough_logging_handler.py | 1 + ...test_batch_embed_content_transformation.py | 36 ++++++++++++++ ...test_vertex_passthrough_logging_handler.py | 47 +++++++++++++++++++ 5 files changed, 95 insertions(+), 1 deletion(-) 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 706e76932e8..654065245ee 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 @@ -288,6 +288,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, _predictions=_predictions, input=input, + raw_usage_metadata=_json_response.get("usageMetadata"), ) async def async_batch_embeddings( @@ -395,4 +396,5 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, _predictions=_predictions, input=input, + raw_usage_metadata=_json_response.get("usageMetadata"), ) 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 96a785c06c9..af78b883b94 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 @@ -486,8 +486,9 @@ _usage_metadata_adapter: Final = TypeAdapter(UsageMetadata) def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None: if not isinstance(raw_usage_metadata, dict): return None + details: Final = raw_usage_metadata.get("promptTokensDetails", raw_usage_metadata.get("promptTokenDetails")) try: - return _usage_metadata_adapter.validate_python(raw_usage_metadata) + return _usage_metadata_adapter.validate_python({**raw_usage_metadata, "promptTokensDetails": details or []}) except ValidationError: return None @@ -615,6 +616,7 @@ def process_response( model_response: EmbeddingResponse, model: str, _predictions: VertexAIBatchEmbeddingsResponseObject, + raw_usage_metadata: object = None, ) -> EmbeddingResponse: openai_embeddings: Final[list[Embedding]] = [] for idx, embedding in enumerate(_predictions["embeddings"]): @@ -628,6 +630,12 @@ def process_response( model_response.data = openai_embeddings model_response.model = model + if _parse_usage_metadata(raw_usage_metadata) is not None: + model_response.usage = _usage_from_embed_content_response( + input=input, model=model, raw_usage_metadata=raw_usage_metadata, resolved_files={} + ) + return model_response + has_nested: Final = isinstance(input, list) and any(isinstance(e, list) for e in input) if _is_multimodal_input(input) or has_nested: input_list: Final = input if isinstance(input, list) else [input] diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index d5c162296d3..2fe99e51a18 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -555,6 +555,7 @@ class VertexPassthroughLoggingHandler: model_response=model_response, model=model, _predictions=response_json, + raw_usage_metadata=response_json.get("usageMetadata"), ) else: litellm_embedding_response = process_embed_content_response( diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index d7d921ab35c..e07399e6dbc 100644 --- a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -652,3 +652,39 @@ def test_gemini_embedding(respx_mock: respx.MockRouter) -> None: expected_prompt_tokens: Final = litellm.token_counter(model="gemini/gemini-embedding-001", text="Hello, world!") assert response.usage.prompt_tokens == expected_prompt_tokens assert response.usage.total_tokens == expected_prompt_tokens + + +@pytest.mark.respx(assert_all_called=True) +def test_gemini_batch_embedding_usage_comes_from_usage_metadata(respx_mock: respx.MockRouter) -> None: + respx_mock.post( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2:batchEmbedContents" + ).mock( + return_value=httpx.Response( + 200, + json={ + "embeddings": [{"values": [0.1, 0.2]}, {"values": [0.3, 0.4]}], + "usageMetadata": { + "promptTokenCount": 262, + "promptTokenDetails": [ + {"modality": "TEXT", "tokenCount": 4}, + {"modality": "IMAGE", "tokenCount": 258}, + ], + }, + }, + ) + ) + + response: Final = litellm.embedding( + model="gemini/gemini-embedding-2", + input=["a red apple", IMAGE_DATA_URI], + api_key="gemini-test-key", + ) + + assert response.usage.prompt_tokens == 262 + assert response.usage.total_tokens == 262 + assert response.usage.prompt_tokens_details.text_tokens == 4 + assert response.usage.prompt_tokens_details.image_tokens == 258 + prices: Final = litellm.model_cost["gemini/gemini-embedding-2"] + assert litellm.completion_cost(completion_response=response) == pytest.approx( + 4 * prices["input_cost_per_token"] + 258 * prices["input_cost_per_image_token"] + ) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py index d05ab8a9bcf..fa6a1ef2347 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py @@ -486,3 +486,50 @@ def test_build_complete_streaming_response_assembles_stream_generate_content_sse assert response.choices[0].message.content == "Hello there!" assert response.choices[0].finish_reason == "stop" assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (10, 8, 18) + + +def test_gemini_batch_embed_contents_passthrough_bills_image_tokens_from_usage_metadata(): + logging_obj: Final = Logging( + model="unknown", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime(2026, 1, 1), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.optional_params = {} + response: Final = httpx.Response( + 200, + json={ + "embeddings": [{"values": [0.1, 0.2]}], + "usageMetadata": { + "promptTokenCount": 258, + "promptTokenDetails": [{"modality": "IMAGE", "tokenCount": 258}], + }, + }, + ) + + result: Final = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2:batchEmbedContents", + result=response.text, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + cache_hit=False, + request_body={ + "requests": [ + { + "model": "models/gemini-embedding-2", + "content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "iVBORw0KGgo="}}]}, + } + ] + }, + ) + + assert result["result"].usage.prompt_tokens == 258 + assert result["result"].usage.prompt_tokens_details.image_tokens == 258 + assert result["kwargs"]["response_cost"] == pytest.approx( + 258 * litellm.model_cost["gemini/gemini-embedding-2"]["input_cost_per_image_token"] + )