scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

fix(bedrock): populate completion_tokens_details in converse _transform_usage

Browse source 
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sameer Kankute 2026-03-10 13:17:28 +05:30
parent 3e1f343d1a
commit 0122abfa2b
2 changed files with 48 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -51,6 +51,7 @@ from litellm.types.llms.openai import (
)
from litellm.types.utils import (
ChatCompletionMessageToolCall,
CompletionTokensDetailsWrapper,
Function,
Message,
ModelResponse,
@ -63,6 +64,7 @@ from litellm.utils import (
has_tool_call_blocks,
last_assistant_with_tool_calls_has_no_thinking_blocks,
supports_reasoning,
token_counter,
)
from ..common_utils import (
@ -1621,7 +1623,11 @@ class AmazonConverseConfig(BaseConfig):
thinking_blocks_list.append(_redacted_block)
return thinking_blocks_list
def _transform_usage(self, usage: ConverseTokenUsageBlock) -> Usage:
def _transform_usage(
self,
usage: ConverseTokenUsageBlock,
reasoning_content: Optional[str] = None,
) -> Usage:
input_tokens = usage["inputTokens"]
output_tokens = usage["outputTokens"]
total_tokens = usage["totalTokens"]
@ -1638,6 +1644,19 @@ class AmazonConverseConfig(BaseConfig):
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens
)
reasoning_tokens = (
token_counter(text=reasoning_content, count_response_tokens=True)
if reasoning_content
else 0
)
completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens,
text_tokens=(
output_tokens - reasoning_tokens
if reasoning_tokens > 0
else output_tokens
),
)
openai_usage = Usage(
prompt_tokens=input_tokens,
completion_tokens=output_tokens,
@ -1645,6 +1664,7 @@ class AmazonConverseConfig(BaseConfig):
prompt_tokens_details=prompt_tokens_details,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
completion_tokens_details=completion_tokens_details,
)
return openai_usage
@ -1981,7 +2001,10 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message["tool_calls"] = filtered_tools
## CALCULATING USAGE - bedrock returns usage in the headers
usage = self._transform_usage(completion_response["usage"])
usage = self._transform_usage(
completion_response["usage"],
reasoning_content=chat_completion_message.get("reasoning_content"),
)
## HANDLE TOOL CALLS
_message = Message(**chat_completion_message)

23
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -43,6 +43,29 @@ def test_transform_usage():
)
assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"]
assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"]
# completion_tokens_details should always be populated
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 == usage["outputTokens"]
def test_transform_usage_with_reasoning_content():
"""Test that completion_tokens_details correctly tracks reasoning vs text tokens."""
usage = ConverseTokenUsageBlock(
**{
"inputTokens": 10,
"outputTokens": 100,
"totalTokens": 110,
}
)
config = AmazonConverseConfig()
reasoning_text = "Let me think about this step by step."
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 == (
usage["outputTokens"] - openai_usage.completion_tokens_details.reasoning_tokens
)
def test_transform_system_message():