mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(bedrock): parse cacheDetails for Converse 1h/5m cache write cost split
AmazonConverseConfig._transform_usage only read the aggregate cacheWriteInputTokens field, so cache_creation_token_details was always unset for Bedrock Converse responses. calculate_cache_writing_cost bills the whole cache-write count at the 5m rate whenever that field is None, so 1-hour TTL cache writes on the standard Bedrock chat path were always undercounted, even though Bedrock returns the 5m/1h split in usage.cacheDetails. Parse cacheDetails (when present) into CacheCreationTokenDetails so the correct rate applies to each portion. No cacheDetails in the response (older models/regions) keeps the previous behavior. Fixes #36760 Co-Authored-By: pi (Claude/GPT via @earendil-works/pi-coding-agent) <noreply@earendil.works>
This commit is contained in:
parent
a7397b2459
commit
97290b4e0e
3 changed files with 74 additions and 3 deletions
|
|
@ -57,6 +57,7 @@ from litellm.types.llms.openai import (
|
|||
OpenAIMessageContentListBlock,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CacheCreationTokenDetails,
|
||||
ChatCompletionMessageToolCall,
|
||||
CompletionTokensDetailsWrapper,
|
||||
Function,
|
||||
|
|
@ -1770,6 +1771,24 @@ class AmazonConverseConfig(BaseConfig):
|
|||
thinking_blocks_list.append(_redacted_block)
|
||||
return thinking_blocks_list
|
||||
|
||||
@staticmethod
|
||||
def _parse_cache_details(usage: ConverseTokenUsageBlock) -> "CacheCreationTokenDetails | None":
|
||||
"""
|
||||
Split Converse's aggregate cacheWriteInputTokens into the 5m/1h TTL
|
||||
breakdown from `cacheDetails`, so cost calc can bill each tier
|
||||
correctly instead of defaulting the whole write to the 5m rate.
|
||||
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html
|
||||
"""
|
||||
cache_details = usage.get("cacheDetails")
|
||||
if not cache_details:
|
||||
return None
|
||||
tokens_5m = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m")
|
||||
tokens_1h = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h")
|
||||
return CacheCreationTokenDetails(
|
||||
ephemeral_5m_input_tokens=tokens_5m,
|
||||
ephemeral_1h_input_tokens=tokens_1h,
|
||||
)
|
||||
|
||||
def _transform_usage(
|
||||
self,
|
||||
usage: ConverseTokenUsageBlock,
|
||||
|
|
@ -1792,6 +1811,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
prompt_tokens_details: Final = PromptTokensDetailsWrapper(
|
||||
cached_tokens=cache_read_input_tokens,
|
||||
cache_creation_tokens=cache_creation_input_tokens,
|
||||
cache_creation_token_details=self._parse_cache_details(usage),
|
||||
text_tokens=raw_input_tokens,
|
||||
)
|
||||
reasoning_tokens = token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
|
||||
|
|
|
|||
|
|
@ -216,14 +216,22 @@ class ConverseResponseOutputBlock(TypedDict):
|
|||
message: MessageBlock | None
|
||||
|
||||
|
||||
class ConverseTokenUsageBlock(TypedDict):
|
||||
class CacheDetailBlock(TypedDict):
|
||||
"""Per-TTL cache-write breakdown. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
|
||||
|
||||
inputTokens: int
|
||||
outputTokens: int
|
||||
totalTokens: int
|
||||
ttl: Literal["5m", "1h"]
|
||||
|
||||
|
||||
class ConverseTokenUsageBlock(TypedDict, total=False):
|
||||
inputTokens: Required[int]
|
||||
outputTokens: Required[int]
|
||||
totalTokens: Required[int]
|
||||
cacheReadInputTokenCount: int
|
||||
cacheReadInputTokens: int
|
||||
cacheWriteInputTokenCount: int
|
||||
cacheWriteInputTokens: int
|
||||
cacheDetails: list[CacheDetailBlock]
|
||||
|
||||
|
||||
class ServiceTierBlock(TypedDict):
|
||||
|
|
|
|||
|
|
@ -51,6 +51,49 @@ def test_transform_usage():
|
|||
assert openai_usage.completion_tokens_details.text_tokens == usage["outputTokens"]
|
||||
|
||||
|
||||
def test_transform_usage_with_cache_details():
|
||||
"""cacheDetails should split cacheWriteInputTokens into the 5m/1h TTL breakdown
|
||||
so cost calc can bill the 1h portion at its own (higher) rate instead of
|
||||
defaulting the whole write to the 5m rate. See issue #36760."""
|
||||
usage = ConverseTokenUsageBlock(
|
||||
**{
|
||||
"inputTokens": 76,
|
||||
"outputTokens": 259,
|
||||
"totalTokens": 335,
|
||||
"cacheWriteInputTokens": 362,
|
||||
"cacheDetails": [
|
||||
{"inputTokens": 74, "ttl": "1h"},
|
||||
{"inputTokens": 288, "ttl": "5m"},
|
||||
],
|
||||
}
|
||||
)
|
||||
config = AmazonConverseConfig()
|
||||
openai_usage = config._transform_usage(usage)
|
||||
details = openai_usage.prompt_tokens_details.cache_creation_token_details
|
||||
assert details is not None
|
||||
assert details.ephemeral_1h_input_tokens == 74
|
||||
assert details.ephemeral_5m_input_tokens == 288
|
||||
|
||||
|
||||
def test_transform_usage_without_cache_details_stays_none():
|
||||
"""No cacheDetails in the response (older models/regions) should leave
|
||||
cache_creation_token_details unset, same as before this field existed."""
|
||||
usage = ConverseTokenUsageBlock(
|
||||
**{
|
||||
"inputTokens": 3,
|
||||
"outputTokens": 401,
|
||||
"totalTokens": 2193,
|
||||
"cacheWriteInputTokens": 1789,
|
||||
}
|
||||
)
|
||||
config = AmazonConverseConfig()
|
||||
openai_usage = config._transform_usage(usage)
|
||||
assert (
|
||||
getattr(openai_usage.prompt_tokens_details, "cache_creation_token_details", None)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_transform_usage_with_reasoning_content():
|
||||
"""Test that completion_tokens_details correctly tracks reasoning vs text tokens."""
|
||||
usage = ConverseTokenUsageBlock(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue