diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 06010c706e3..4da0b0d888f 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1824 + "limit": 1823 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 9681d64f656..8a7ebfb7949 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -424,6 +424,38 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int: return 0 +def _is_converse_usage_shape(usage_object: dict) -> bool: + """Converse-family models report camelCase token counts, not Anthropic's snake_case.""" + return "inputTokens" in usage_object or "outputTokens" in usage_object + + +def _get_converse_batch_usage(usage_object: dict) -> Usage: + """Read a Converse-shaped usage block with the same transform the live Converse path uses, + so a batch and an equivalent non-batch call agree on tokens.""" + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.llms.bedrock import ConverseTokenUsageBlock + + input_tokens: Final = int(usage_object.get("inputTokens") or 0) + output_tokens: Final = int(usage_object.get("outputTokens") or 0) + cache_read: Final = int( + usage_object.get("cacheReadInputTokens") or usage_object.get("cacheReadInputTokenCount") or 0 + ) + cache_write: Final = int( + usage_object.get("cacheWriteInputTokens") or usage_object.get("cacheWriteInputTokenCount") or 0 + ) + return AmazonConverseConfig().transform_usage( + ConverseTokenUsageBlock( + inputTokens=input_tokens, + outputTokens=output_tokens, + totalTokens=int(usage_object.get("totalTokens") or input_tokens + output_tokens), + cacheReadInputTokenCount=cache_read, + cacheReadInputTokens=cache_read, + cacheWriteInputTokenCount=cache_write, + cacheWriteInputTokens=cache_write, + ) + ) + + def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_provider: str = "openai") -> Usage: """ Get the tokens of a batch job from the response body @@ -431,10 +463,21 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov if custom_llm_provider in ("anthropic", "bedrock"): from litellm.llms.anthropic.chat.transformation import AnthropicConfig - return AnthropicConfig().calculate_usage( - usage_object=response_body.get("usage", None) or {}, + usage_object: Final = response_body.get("usage", None) or {} + if custom_llm_provider == "bedrock" and _is_converse_usage_shape(usage_object): + return _get_converse_batch_usage(usage_object) + anthropic_usage: Final = AnthropicConfig().calculate_usage( + usage_object=usage_object, reasoning_content=None, ) + if usage_object and anthropic_usage.total_tokens == 0: + verbose_logger.warning( + "batch output line reported usage this parser does not understand, so it will be billed at $0. " + "provider=%s usage_keys=%s", + custom_llm_provider, + sorted(usage_object.keys()), + ) + return anthropic_usage from litellm.responses.utils import ResponseAPILoggingUtils _usage_dict: Final = response_body.get("usage", None) or {} diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 85918d40e12..d7c6725671a 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1770,7 +1770,7 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list - def _transform_usage( + def transform_usage( self, usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, @@ -2191,7 +2191,7 @@ class AmazonConverseConfig(BaseConfig): chat_completion_message["tool_calls"] = filtered_tools ## CALCULATING USAGE - bedrock returns usage in the headers - usage: Final = self._transform_usage( + usage: Final = self.transform_usage( completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), ) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 8d2b3dae71b..86f7e9b0d9f 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -559,7 +559,7 @@ class AWSEventStreamDecoder: elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) elif "usage" in chunk_data: - usage = converse_config._transform_usage(chunk_data.get("usage", {})) + usage = converse_config.transform_usage(chunk_data.get("usage", {})) model_response_provider_specific_fields: Final = {} if "trace" in chunk_data: diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index d2074853f2b..84afee44e92 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -15,6 +15,7 @@ deterministic stand-ins so the arithmetic under test is the only variable. """ import json +import logging import os import sys from types import MappingProxyType @@ -1299,3 +1300,60 @@ async def test_output_file_content_bedrock_reads_with_deployment_aws_credentials assert captured["aws_region_name"] == "us-west-2" assert captured["_litellm_internal_model_credentials"] is snapshot assert "model" not in captured + + +# =========================================================================== # +# Bedrock batch usage is parsed by the shape of the payload, not the provider +# +# Regression: every bedrock batch line went through the Anthropic usage parser, +# which reads snake_case input_tokens/output_tokens. A Converse-family model +# (Nova and friends) reports camelCase inputTokens/outputTokens, so usage read +# 0/0/0 and the batch billed $0 despite real token consumption. +# =========================================================================== # + + +def test_bedrock_converse_shaped_batch_usage_is_parsed(): + body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}} + usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742) + + +def test_bedrock_converse_batch_usage_totals_default_when_absent(): + body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 10, "outputTokens": 4}} + usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 4, 14) + + +def test_bedrock_converse_batch_usage_includes_cache_tokens(): + body = { + "model": "us.amazon.nova-lite-v1:0", + "usage": { + "inputTokens": 100, + "outputTokens": 20, + "totalTokens": 120, + "cacheReadInputTokens": 800, + "cacheWriteInputTokens": 200, + }, + } + usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") + assert usage.prompt_tokens == 1100 + assert usage.completion_tokens == 20 + assert usage.prompt_tokens_details.cached_tokens == 800 + assert usage.prompt_tokens_details.cache_creation_tokens == 200 + + +def test_bedrock_anthropic_shaped_batch_usage_still_parsed(): + """Anthropic-shaped bedrock output (what an Anthropic model's batch emits) must not regress.""" + body = {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 18, "output_tokens": 10}} + usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (18, 10, 28) + + +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}} + with caplog.at_level(logging.WARNING): + usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") + assert usage.total_tokens == 0 + assert "does not understand" in caplog.text + assert "inputTextTokenCount" in caplog.text diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index d1d1f9ab489..ee6db582c46 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -30,7 +30,7 @@ def test_transform_usage(): } ) config = AmazonConverseConfig() - openai_usage = config._transform_usage(usage) + openai_usage = config.transform_usage(usage) assert ( openai_usage.prompt_tokens == usage["inputTokens"] @@ -62,7 +62,7 @@ def test_transform_usage_with_reasoning_content(): ) config = AmazonConverseConfig() reasoning_text = "Let me think about this step by step." - openai_usage = config._transform_usage(usage, reasoning_content=reasoning_text) + openai_usage = config.transform_usage(usage, reasoning_content=reasoning_text) assert openai_usage.completion_tokens_details is not None assert openai_usage.completion_tokens_details.reasoning_tokens > 0 assert openai_usage.completion_tokens_details.text_tokens == (