Merge c6616a3e2f into 72228db5c6

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"),
             )

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]

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(

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"]
+    )

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"]
+    )