Merge pull request #41767 from BerriAI/litellm_bedrock_titan_batch_usage

fix(batches): bill Bedrock Titan embedding batch lines from inputTextTokenCount

litellm/batches/batch_utils.py

@@ -9,6 +9,7 @@ import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
+from litellm.llms.bedrock.batches.transformation import titan_embedding_usage_from_batch_output
 from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details
 from litellm.types.llms.openai import Batch
 from litellm.types.utils import ModelInfo, Usage
@@ -673,6 +674,11 @@ def _get_batch_job_usage_from_response_body(
         from litellm.llms.anthropic.chat.transformation import AnthropicConfig
         from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
 
+        titan_usage: Final = (
+            titan_embedding_usage_from_batch_output(response_body) if custom_llm_provider == "bedrock" else None
+        )
+        if titan_usage is not None:
+            return titan_usage
         usage_object: Final = response_body.get("usage", None) or {}
         if custom_llm_provider == "bedrock" and AmazonConverseConfig.is_converse_usage_shape(usage_object):
             return AmazonConverseConfig().usage_from_batch_output(usage_object)

litellm/llms/bedrock/batches/transformation.py

@@ -1,6 +1,7 @@
 import os
 import re
 import time
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 
 from httpx import Headers, Response
@@ -26,7 +27,7 @@ from litellm.types.llms.openai import (
     AllMessageValues,
     CreateBatchRequest,
 )
-from litellm.types.utils import LiteLLMBatch, LlmProviders
+from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage
 
 from ..base_aws_llm import BaseAWSLLM
 from ..common_utils import (
@@ -60,6 +61,20 @@ def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]:
         ) from e
 
 
+def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None:
+    """Titan embedding batch lines report usage as a top-level inputTextTokenCount, not a usage block."""
+    if "embedding" not in model_output and "embeddingsByType" not in model_output:
+        return None
+    input_text_token_count: Final = model_output.get("inputTextTokenCount")
+    if isinstance(input_text_token_count, bool) or not isinstance(input_text_token_count, int):
+        return None
+    return Usage(
+        prompt_tokens=input_text_token_count,
+        completion_tokens=0,
+        total_tokens=input_text_token_count,
+    )
+
+
 class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
     """
     Config for Bedrock Batches - handles batch job creation and management for Bedrock

tests/test_litellm/batches/test_batch_utils.py

@@ -1755,6 +1755,44 @@ def test_bedrock_anthropic_shaped_batch_usage_still_parsed():
     assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (18, 10, 28)
 
 
+def test_bedrock_titan_embedding_batch_usage_is_parsed():
+    """Titan embedding batch lines carry a top-level inputTextTokenCount and no usage block."""
+    body = {"embedding": [0.1, 0.2], "embeddingsByType": {"float": [0.1, 0.2]}, "inputTextTokenCount": 17}
+    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
+    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (17, 0, 17)
+
+
+def test_bedrock_titan_embedding_batch_is_billed():
+    """Binary embedding rows carry only embeddingsByType and must bill like float rows."""
+    rows = [
+        {"recordId": "0", "modelOutput": {"embedding": [0.1], "inputTextTokenCount": 10}},
+        {"recordId": "1", "modelOutput": {"embeddingsByType": {"binary": [1, 0]}, "inputTextTokenCount": 7}},
+    ]
+    result = bu._aggregate_batch_cost_usage_models(
+        entries=rows,
+        custom_llm_provider="bedrock",
+        model_name="amazon.titan-embed-text-v2:0",
+        model_info={"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 0.0},
+    )
+    assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (17, 0, 17)
+    assert result.cost == pytest.approx(17 * 1e-6)
+
+
+@pytest.mark.parametrize(
+    "body",
+    [
+        {"embedding": [0.1], "inputTextTokenCount": "17"},
+        {"embedding": [0.1], "inputTextTokenCount": True},
+        {"embedding": [0.1], "inputTextTokenCount": None},
+        {"results": [{"outputText": "hi", "tokenCount": 2}], "inputTextTokenCount": 17},
+    ],
+)
+def test_bedrock_input_text_token_count_outside_embedding_lines_is_not_billed(body):
+    """Only embedding lines are parsed here; Titan text generation lines are left as they were."""
+    usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
+    assert usage.total_tokens == 0
+
+
 def test_unparsable_bedrock_batch_usage_warns(caplog):
     """An unrecognized usage shape must be visible, not a silent $0."""
     body = {"model": "amazon.titan-text-lite-v1", "usage": {"inputTextTokenCount": 42}}