diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 4f74e3f7035..ae0f8c5935b 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -63,7 +63,7 @@ def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: 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: + 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): diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index a2811864519..9a089112c70 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1763,9 +1763,10 @@ def test_bedrock_titan_embedding_batch_usage_is_parsed(): def test_bedrock_titan_embedding_batch_is_billed(): + """Binary embedding rows carry only embeddingsByType and must bill like float rows.""" rows = [ - {"recordId": str(i), "modelOutput": {"embedding": [0.1], "inputTextTokenCount": count}} - for i, count in enumerate((10, 7)) + {"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,