scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 77

fix(batches): bill Titan binary embedding batch lines that only carry embeddingsByType

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-18 05:32:25 +00:00
parent e50fc8ba75
commit 0247e9b634
2 changed files with 4 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/bedrock/batches/transformation.py
View file

@ -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):

5
tests/test_litellm/batches/test_batch_utils.py
View file

@ -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,