mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(bedrock): guard cache-detail split against partial/unrecognized ttl entries
Address review feedback on #36762: - Only use the parsed 5m/1h split when it fully accounts for cacheWriteInputTokens; an unrecognized ttl or missing entry now falls back to the aggregate (previous behavior) instead of silently understating cost. - Mark TypedDict fields ReadOnly (AWS response data, never constructed by us) to satisfy the repo's type-discipline lint gate. - Trim comments and add Final to locals per repo style. Co-Authored-By: pi (Claude/GPT via @earendil-works/pi-coding-agent) <noreply@earendil.works>
This commit is contained in:
parent
97290b4e0e
commit
42a2b5f057
3 changed files with 42 additions and 21 deletions
|
|
@ -1773,17 +1773,17 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
@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")
|
||||
"""https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
|
||||
cache_details: Final = 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")
|
||||
tokens_5m: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "5m")
|
||||
tokens_1h: Final = sum(d["inputTokens"] for d in cache_details if d.get("ttl") == "1h")
|
||||
# An unrecognized ttl or a partial breakdown would silently understate
|
||||
# the cache-write cost, so only use the split when it fully accounts
|
||||
# for the aggregate; otherwise fall back to the aggregate-only (5m) cost.
|
||||
if tokens_5m + tokens_1h != usage.get("cacheWriteInputTokens", 0):
|
||||
return None
|
||||
return CacheCreationTokenDetails(
|
||||
ephemeral_5m_input_tokens=tokens_5m,
|
||||
ephemeral_1h_input_tokens=tokens_1h,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import json
|
|||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from typing_extensions import Required, TypedDict, override
|
||||
from typing_extensions import ReadOnly, Required, TypedDict, override
|
||||
|
||||
from .openai import ChatCompletionToolCallChunk
|
||||
|
||||
|
|
@ -217,21 +217,21 @@ class ConverseResponseOutputBlock(TypedDict):
|
|||
|
||||
|
||||
class CacheDetailBlock(TypedDict):
|
||||
"""Per-TTL cache-write breakdown. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
|
||||
"""Per-TTL cache-write breakdown, read-only AWS response data. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CacheDetail.html"""
|
||||
|
||||
inputTokens: int
|
||||
ttl: Literal["5m", "1h"]
|
||||
inputTokens: ReadOnly[int]
|
||||
ttl: ReadOnly[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]
|
||||
inputTokens: Required[ReadOnly[int]]
|
||||
outputTokens: Required[ReadOnly[int]]
|
||||
totalTokens: Required[ReadOnly[int]]
|
||||
cacheReadInputTokenCount: ReadOnly[int]
|
||||
cacheReadInputTokens: ReadOnly[int]
|
||||
cacheWriteInputTokenCount: ReadOnly[int]
|
||||
cacheWriteInputTokens: ReadOnly[int]
|
||||
cacheDetails: ReadOnly[list[CacheDetailBlock]] # mutable-ok: AWS response array, never mutated after parsing
|
||||
|
||||
|
||||
class ServiceTierBlock(TypedDict):
|
||||
|
|
|
|||
|
|
@ -75,6 +75,27 @@ def test_transform_usage_with_cache_details():
|
|||
assert details.ephemeral_5m_input_tokens == 288
|
||||
|
||||
|
||||
def test_transform_usage_with_mismatched_cache_details_falls_back():
|
||||
"""An unrecognized ttl or partial breakdown must not silently understate
|
||||
cache-write cost, so the split is only used when it fully accounts for
|
||||
cacheWriteInputTokens."""
|
||||
usage = ConverseTokenUsageBlock(
|
||||
**{
|
||||
"inputTokens": 76,
|
||||
"outputTokens": 259,
|
||||
"totalTokens": 335,
|
||||
"cacheWriteInputTokens": 362,
|
||||
"cacheDetails": [{"inputTokens": 74, "ttl": "1h"}], # missing the 5m entry
|
||||
}
|
||||
)
|
||||
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_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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue