fix(bedrock): make prompt caching work on the Nova InvokeModel route

Nova InvokeModel rejects the standalone cachePoint blocks the shared Converse transform emits, so each one is folded into the block it caches and tool_config injection points are dropped before the transform runs, since this route has no tool caching to credit. Usage reads Bedrock's Count-suffixed cache keys and adds cached tokens into prompt_tokens, streaming routes every wrapped InvokeModel event through the Converse chunk parser and tolerates the missing totalTokens, and the Nova 1 cost-map entries gain cache_read_input_token_cost at a quarter of the input rate
This commit is contained in:
mateo-berri 2026-09-15 18:01:21 -07:00
parent 226b1e1bb9
commit 3c15f64fd4
9 changed files with 367 additions and 35 deletions

View file

@ -1902,7 +1902,7 @@ class AmazonConverseConfig(BaseConfig):
return None
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")
if tokens_5m + tokens_1h != usage.get("cacheWriteInputTokens", 0):
if tokens_5m + tokens_1h != AmazonConverseConfig._cache_write_count(usage):
return None
return CacheCreationTokenDetails(
ephemeral_5m_input_tokens=tokens_5m,
@ -1933,6 +1933,15 @@ class AmazonConverseConfig(BaseConfig):
return int(value)
return 0
@staticmethod
def _cache_read_count(usage_object: Mapping[str, object]) -> int:
"""Converse reports ``cacheReadInputTokens``; InvokeModel reports ``cacheReadInputTokenCount``."""
return AmazonConverseConfig._usage_count(usage_object, "cacheReadInputTokens", "cacheReadInputTokenCount")
@staticmethod
def _cache_write_count(usage_object: Mapping[str, object]) -> int:
return AmazonConverseConfig._usage_count(usage_object, "cacheWriteInputTokens", "cacheWriteInputTokenCount")
def usage_from_batch_output(self, usage_object: Mapping[str, object]) -> Usage:
"""Read a Converse-shaped usage block out of a batch output line.
@ -1942,8 +1951,8 @@ class AmazonConverseConfig(BaseConfig):
"""
input_tokens: Final = self._usage_count(usage_object, "inputTokens")
output_tokens: Final = self._usage_count(usage_object, "outputTokens")
cache_read: Final = self._usage_count(usage_object, "cacheReadInputTokens", "cacheReadInputTokenCount")
cache_write: Final = self._usage_count(usage_object, "cacheWriteInputTokens", "cacheWriteInputTokenCount")
cache_read: Final = self._cache_read_count(usage_object)
cache_write: Final = self._cache_write_count(usage_object)
return self.transform_usage(
ConverseTokenUsageBlock(
inputTokens=input_tokens,
@ -1963,19 +1972,12 @@ class AmazonConverseConfig(BaseConfig):
thinking_ran: bool = False,
provider_reasoning_tokens: int | None = None,
) -> Usage:
input_tokens = usage["inputTokens"]
raw_input_tokens: Final = usage["inputTokens"]
output_tokens: Final = usage["outputTokens"]
total_tokens: Final = usage["totalTokens"]
cache_creation_input_tokens: int = 0
cache_read_input_tokens: int = 0
raw_input_tokens: Final = input_tokens # capture before inflation
if "cacheReadInputTokens" in usage:
cache_read_input_tokens = usage["cacheReadInputTokens"]
input_tokens += cache_read_input_tokens
if "cacheWriteInputTokens" in usage:
cache_creation_input_tokens = usage["cacheWriteInputTokens"]
input_tokens += cache_creation_input_tokens
cache_read_input_tokens: Final = self._cache_read_count(usage)
cache_creation_input_tokens: Final = self._cache_write_count(usage)
input_tokens: Final = raw_input_tokens + cache_read_input_tokens + cache_creation_input_tokens
total_tokens: Final = usage.get("totalTokens", input_tokens + output_tokens)
prompt_tokens_details: Final = PromptTokensDetailsWrapper(
cached_tokens=cache_read_input_tokens,

View file

@ -3,6 +3,7 @@ from collections.abc import AsyncIterator, Iterator
from typing import Final, cast
import httpx
from pydantic import TypeAdapter
import litellm
from litellm import verbose_logger
@ -51,6 +52,15 @@ bedrock_tool_name_mappings: Final[InMemoryCache] = InMemoryCache(max_size_in_mem
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
converse_config: Final = AmazonConverseConfig()
NOVA_INVOKE_STREAM_EVENT_TYPES: Final = (
"messageStart",
"contentBlockStart",
"contentBlockDelta",
"contentBlockStop",
"messageStop",
"metadata",
)
NOVA_INVOKE_STREAM_EVENT_PAYLOAD: Final = TypeAdapter(dict[str, object])
class AmazonCohereChatConfig:
@ -601,14 +611,12 @@ class AWSEventStreamDecoder:
if thinking_blocks:
self._thinking_ran = True
carries_message_content: Final = any(
key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason", "trace")
trace: Final = chunk_data.get("trace")
carries_message_content: Final = bool(trace) or any(
key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason")
)
model_response_provider_specific_fields: Final = {}
if "trace" in chunk_data:
trace: Final = chunk_data.get("trace")
model_response_provider_specific_fields["trace"] = trace
model_response_provider_specific_fields: Final = {"trace": trace} if trace else {}
response: Final = ModelResponseStream(
choices=[
StreamingChoices(
@ -654,10 +662,10 @@ class AWSEventStreamDecoder:
):
return self.converse_chunk_parser(chunk_data=chunk_data)
######### /bedrock/invoke nova mappings ###############
elif "contentBlockDelta" in chunk_data:
# when using /bedrock/invoke/nova, the chunk_data is nested under "contentBlockDelta"
_chunk_data: Final = chunk_data.get("contentBlockDelta", {})
return self.converse_chunk_parser(chunk_data=_chunk_data)
elif nova_event_type := next((key for key in NOVA_INVOKE_STREAM_EVENT_TYPES if key in chunk_data), None):
return self.converse_chunk_parser(
chunk_data=NOVA_INVOKE_STREAM_EVENT_PAYLOAD.validate_python(chunk_data[nova_event_type])
)
######## bedrock.mistral mappings ###############
elif "outputs" in chunk_data:
if len(chunk_data["outputs"]) == 1 and chunk_data["outputs"][0].get("text", None) is not None:

View file

@ -6,12 +6,21 @@ Inherits from `AmazonConverseConfig`
Nova + Invoke API Tutorial: https://docs.aws.amazon.com/nova/latest/userguide/using-invoke-api.html
"""
from typing import TYPE_CHECKING, Final
from collections.abc import Callable, Mapping, Sequence
from functools import reduce
from typing import TYPE_CHECKING, Final, TypeVar
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.llms.bedrock import BedrockInvokeNovaRequest
from litellm.types.llms.bedrock import (
BedrockInvokeNovaRequest,
CachePointBlock,
ContentBlock,
MessageBlock,
SystemContentBlock,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
@ -21,6 +30,49 @@ from .base_invoke_transformation import AmazonInvokeConfig
if TYPE_CHECKING:
import tiktoken
_CachePointCarrier = TypeVar("_CachePointCarrier", SystemContentBlock, ContentBlock)
_INJECTION_POINTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def _without_tool_config_injection_points(optional_params: Mapping[str, object]) -> dict[str, object]:
"""InvokeModel has no tool caching, and a ``tool_config`` point the Converse transform
placed would credit the gateway for a cachePoint this request cannot carry.
"""
raw_points: Final = optional_params.get("cache_control_injection_points")
if raw_points is None:
return dict(optional_params)
try:
points = _INJECTION_POINTS.validate_python(raw_points)
except ValidationError:
return dict(optional_params)
return {
**optional_params,
"cache_control_injection_points": [point for point in points if point.get("location") != "tool_config"],
}
def _system_block_with_cache_point(block: SystemContentBlock, cache_point: CachePointBlock) -> SystemContentBlock:
return {**block, "cachePoint": cache_point}
def _content_block_with_cache_point(block: ContentBlock, cache_point: CachePointBlock) -> ContentBlock:
return {**block, "cachePoint": cache_point}
def _inline_cache_points(
blocks: Sequence[_CachePointCarrier],
with_cache_point: Callable[[_CachePointCarrier, CachePointBlock], _CachePointCarrier],
) -> list[_CachePointCarrier]:
def attach(inlined: tuple[_CachePointCarrier, ...], block: _CachePointCarrier) -> tuple[_CachePointCarrier, ...]:
cache_point: Final = block.get("cachePoint")
if cache_point is None or len(block) != 1:
return (*inlined, block)
if not inlined:
return inlined
return (*inlined[:-1], with_cache_point(inlined[-1], cache_point))
return list(reduce(attach, blocks, ()))
class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
"""
@ -46,7 +98,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
self,
model: str,
messages: list[AllMessageValues],
optional_params: dict,
optional_params: dict[str, object],
litellm_params: dict,
headers: dict,
) -> dict:
@ -54,11 +106,13 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
self,
model=model,
messages=messages,
optional_params=optional_params,
optional_params=_without_tool_config_injection_points(optional_params),
litellm_params=litellm_params,
headers=headers,
)
_bedrock_invoke_nova_request: Final = BedrockInvokeNovaRequest(**_transformed_nova_request)
_bedrock_invoke_nova_request: Final = self._inline_cache_points(
BedrockInvokeNovaRequest(**_transformed_nova_request)
)
self._remove_empty_system_messages(_bedrock_invoke_nova_request)
bedrock_invoke_nova_request: Final = self._filter_allowed_fields(_bedrock_invoke_nova_request)
return bedrock_invoke_nova_request
@ -92,6 +146,23 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
json_mode,
)
@staticmethod
def _inline_cache_points(request: BedrockInvokeNovaRequest) -> BedrockInvokeNovaRequest:
"""InvokeModel takes ``cachePoint`` as a key of the block it caches and rejects the
standalone ``{"cachePoint": ...}`` blocks Converse accepts.
"""
return {
**request,
"system": _inline_cache_points(request.get("system", []), _system_block_with_cache_point),
"messages": [
MessageBlock(
role=message["role"],
content=_inline_cache_points(message["content"], _content_block_with_cache_point),
)
for message in request.get("messages", [])
],
}
def _filter_allowed_fields(self, bedrock_invoke_nova_request: BedrockInvokeNovaRequest) -> dict:
"""
Filter out fields that are not allowed in the `BedrockInvokeNovaRequest` dataclass.

View file

@ -353,6 +353,7 @@
"supports_pdf_input": true
},
"amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 6e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -537,6 +538,7 @@
"supports_audio_input": true
},
"amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
"input_cost_per_token": 3.5e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -550,6 +552,7 @@
"supports_tool_choice": true
},
"amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 8e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -2884,6 +2887,7 @@
"supports_function_calling": true
},
"apac.amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.575e-08,
"input_cost_per_token": 6.3e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -2899,6 +2903,7 @@
"supports_tool_choice": true
},
"apac.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 9.25e-09,
"input_cost_per_token": 3.7e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -2912,6 +2917,7 @@
"supports_tool_choice": true
},
"apac.amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.1e-07,
"input_cost_per_token": 8.4e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -12449,6 +12455,7 @@
"source": "https://aws.amazon.com/bedrock/pricing/"
},
"bedrock/us-gov-east-1/amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.4e-07,
"input_cost_per_token": 9.6e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 300000,
@ -12629,6 +12636,7 @@
"supports_audio_input": true
},
"bedrock/us-gov-west-1/amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.8e-08,
"input_cost_per_token": 7.2e-08,
"litellm_provider": "bedrock",
"max_input_tokens": 300000,
@ -12644,6 +12652,7 @@
"supports_tool_choice": true
},
"bedrock/us-gov-west-1/amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 1.05e-08,
"input_cost_per_token": 4.2e-08,
"litellm_provider": "bedrock",
"max_input_tokens": 128000,
@ -12657,6 +12666,7 @@
"supports_tool_choice": true
},
"bedrock/us-gov-west-1/amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.4e-07,
"input_cost_per_token": 9.6e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 300000,
@ -21195,6 +21205,7 @@
"supports_embedding_image_input": true
},
"eu.amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.95e-08,
"input_cost_per_token": 7.8e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -21210,6 +21221,7 @@
"supports_tool_choice": true
},
"eu.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 1.15e-08,
"input_cost_per_token": 4.6e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -21223,6 +21235,7 @@
"supports_tool_choice": true
},
"eu.amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.625e-07,
"input_cost_per_token": 1.05e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -44643,6 +44656,7 @@
"source": "https://aws.amazon.com/polly/pricing/"
},
"us.amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 6e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -44658,6 +44672,7 @@
"supports_tool_choice": true
},
"us.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
"input_cost_per_token": 3.5e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -44686,6 +44701,7 @@
"supports_vision": true
},
"us.amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 8e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,

View file

@ -231,7 +231,7 @@ class CacheDetailBlock(TypedDict):
class ConverseTokenUsageBlock(TypedDict, total=False):
inputTokens: Required[ReadOnly[int]]
outputTokens: Required[ReadOnly[int]]
totalTokens: Required[ReadOnly[int]]
totalTokens: ReadOnly[int]
cacheReadInputTokenCount: ReadOnly[int]
cacheReadInputTokens: ReadOnly[int]
cacheWriteInputTokenCount: ReadOnly[int]

View file

@ -353,6 +353,7 @@
"supports_pdf_input": true
},
"amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 6e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -537,6 +538,7 @@
"supports_audio_input": true
},
"amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
"input_cost_per_token": 3.5e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -550,6 +552,7 @@
"supports_tool_choice": true
},
"amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 8e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -2884,6 +2887,7 @@
"supports_function_calling": true
},
"apac.amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.575e-08,
"input_cost_per_token": 6.3e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -2899,6 +2903,7 @@
"supports_tool_choice": true
},
"apac.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 9.25e-09,
"input_cost_per_token": 3.7e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -2912,6 +2917,7 @@
"supports_tool_choice": true
},
"apac.amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.1e-07,
"input_cost_per_token": 8.4e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -12449,6 +12455,7 @@
"source": "https://aws.amazon.com/bedrock/pricing/"
},
"bedrock/us-gov-east-1/amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.4e-07,
"input_cost_per_token": 9.6e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 300000,
@ -12629,6 +12636,7 @@
"supports_audio_input": true
},
"bedrock/us-gov-west-1/amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.8e-08,
"input_cost_per_token": 7.2e-08,
"litellm_provider": "bedrock",
"max_input_tokens": 300000,
@ -12644,6 +12652,7 @@
"supports_tool_choice": true
},
"bedrock/us-gov-west-1/amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 1.05e-08,
"input_cost_per_token": 4.2e-08,
"litellm_provider": "bedrock",
"max_input_tokens": 128000,
@ -12657,6 +12666,7 @@
"supports_tool_choice": true
},
"bedrock/us-gov-west-1/amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.4e-07,
"input_cost_per_token": 9.6e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 300000,
@ -21195,6 +21205,7 @@
"supports_embedding_image_input": true
},
"eu.amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.95e-08,
"input_cost_per_token": 7.8e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -21210,6 +21221,7 @@
"supports_tool_choice": true
},
"eu.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 1.15e-08,
"input_cost_per_token": 4.6e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -21223,6 +21235,7 @@
"supports_tool_choice": true
},
"eu.amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2.625e-07,
"input_cost_per_token": 1.05e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -44643,6 +44656,7 @@
"source": "https://aws.amazon.com/polly/pricing/"
},
"us.amazon.nova-lite-v1:0": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 6e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,
@ -44658,6 +44672,7 @@
"supports_tool_choice": true
},
"us.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
"input_cost_per_token": 3.5e-08,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 128000,
@ -44686,6 +44701,7 @@
"supports_vision": true
},
"us.amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 8e-07,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 300000,

View file

@ -0,0 +1,67 @@
import json
from litellm.llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import (
AmazonInvokeNovaConfig,
)
from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY
MODEL = "us.amazon.nova-pro-v1:0"
EPHEMERAL = {"type": "ephemeral"}
DEFAULT_CACHE_POINT = {"type": "default"}
TOOLS = [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}]
def _transform_request(messages, optional_params, litellm_params=None):
return AmazonInvokeNovaConfig().transform_request(
model=MODEL,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params if litellm_params is not None else {},
headers={},
)
def test_cache_points_are_inlined_into_the_block_they_cache(local_model_cost_map):
"""InvokeModel rejects the standalone ``{"cachePoint": ...}`` block Converse emits
(``#/system/1: required key [text] not found``); it wants ``cachePoint`` as a key of the
block being cached."""
request = _transform_request(
messages=[
{"role": "system", "content": [{"type": "text", "text": "long system prompt", "cache_control": EPHEMERAL}]},
{"role": "user", "content": [{"type": "text", "text": "hello", "cache_control": EPHEMERAL}]},
{"role": "assistant", "content": "hi there", "cache_control": EPHEMERAL},
{"role": "user", "content": "again"},
],
optional_params={"max_tokens": 20},
)
assert request["system"] == [{"text": "long system prompt", "cachePoint": DEFAULT_CACHE_POINT}]
assert [message["content"] for message in request["messages"]] == [
[{"text": "hello", "cachePoint": DEFAULT_CACHE_POINT}],
[{"text": "hi there", "cachePoint": DEFAULT_CACHE_POINT}],
[{"text": "again"}],
]
def test_cache_point_with_nothing_before_it_is_dropped():
request = AmazonInvokeNovaConfig._inline_cache_points(
{
"system": [{"cachePoint": DEFAULT_CACHE_POINT}],
"messages": [{"role": "user", "content": [{"cachePoint": DEFAULT_CACHE_POINT}, {"text": "hi"}]}],
}
)
assert request["system"] == []
assert request["messages"] == [{"role": "user", "content": [{"text": "hi"}]}]
def test_tool_config_injection_point_is_neither_placed_nor_credited(local_model_cost_map):
"""InvokeModel has no tool caching, so the point cannot land and the gateway must not be
credited for it in spend attribution."""
metadata = {"user_api_key": "sk-test"}
request = _transform_request(
messages=[{"role": "user", "content": "hi"}],
optional_params={"tools": TOOLS, "cache_control_injection_points": [{"location": "tool_config"}]},
litellm_params={"metadata": metadata, "litellm_metadata": None, "model_info": {"id": "dep-bedrock"}},
)
assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["f"]
assert "cachePoint" not in json.dumps(request)
assert GATEWAY_INJECTED_CACHE_METADATA_KEY not in metadata

View file

@ -139,6 +139,118 @@ def test_bedrock_converse_1h_cache_write_billed_at_1h_rate(monkeypatch):
assert completion_cost == pytest.approx(4 * model_info["output_cost_per_token"])
@pytest.mark.parametrize(
"usage, expected_prompt_tokens, expected_cached_tokens, expected_cache_creation_tokens",
[
pytest.param(
{
"inputTokens": 5,
"outputTokens": 3,
"totalTokens": 12270,
"cacheReadInputTokenCount": 12262,
"cacheWriteInputTokenCount": 0,
},
12267,
12262,
0,
id="invoke-model-cache-read",
),
pytest.param(
{
"inputTokens": 5,
"outputTokens": 3,
"totalTokens": 12270,
"cacheReadInputTokenCount": 0,
"cacheWriteInputTokenCount": 12262,
},
12267,
0,
12262,
id="invoke-model-cache-write",
),
pytest.param(
{
"inputTokens": 5,
"outputTokens": 3,
"cacheReadInputTokenCount": 12262,
"cacheWriteInputTokenCount": 0,
},
12267,
12262,
0,
id="invoke-model-streaming-metadata-without-totalTokens",
),
],
)
def test_transform_usage_reads_invoke_model_count_suffixed_cache_keys(
usage, expected_prompt_tokens, expected_cached_tokens, expected_cache_creation_tokens
):
"""InvokeModel Nova reports ``cacheReadInputTokenCount`` and ``cacheWriteInputTokenCount``
where Converse reports the un-suffixed keys, and ``inputTokens`` excludes both."""
openai_usage = AmazonConverseConfig().transform_usage(ConverseTokenUsageBlock(**usage))
assert openai_usage.prompt_tokens == expected_prompt_tokens
assert openai_usage.prompt_tokens_details.cached_tokens == expected_cached_tokens
assert openai_usage._cache_read_input_tokens == expected_cached_tokens
assert openai_usage._cache_creation_input_tokens == expected_cache_creation_tokens
assert openai_usage.completion_tokens == 3
assert openai_usage.total_tokens == 12270
def test_bedrock_invoke_nova_cache_read_billed_at_discounted_rate(monkeypatch):
"""Nova cache reads are billed at 25% of the input rate; without a
``cache_read_input_token_cost`` entry the cached tokens were billed at nothing."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
usage = ConverseTokenUsageBlock(
**{
"inputTokens": 5,
"outputTokens": 3,
"totalTokens": 12270,
"cacheReadInputTokenCount": 12262,
"cacheWriteInputTokenCount": 0,
}
)
openai_usage = AmazonConverseConfig().transform_usage(usage)
model = "bedrock/invoke/us.amazon.nova-pro-v1:0"
prompt_cost, completion_cost = litellm.cost_calculator.cost_per_token(model=model, usage_object=openai_usage)
model_info = litellm.get_model_info(model=model)
assert model_info["cache_read_input_token_cost"] == pytest.approx(model_info["input_cost_per_token"] * 0.25)
assert prompt_cost == pytest.approx(
5 * model_info["input_cost_per_token"] + 12262 * model_info["cache_read_input_token_cost"]
)
assert prompt_cost > 5 * model_info["input_cost_per_token"]
assert completion_cost == pytest.approx(3 * model_info["output_cost_per_token"])
@pytest.mark.parametrize(
"model",
[
"amazon.nova-micro-v1:0",
"amazon.nova-lite-v1:0",
"amazon.nova-pro-v1:0",
"us.amazon.nova-micro-v1:0",
"us.amazon.nova-lite-v1:0",
"us.amazon.nova-pro-v1:0",
"eu.amazon.nova-micro-v1:0",
"eu.amazon.nova-lite-v1:0",
"eu.amazon.nova-pro-v1:0",
"apac.amazon.nova-micro-v1:0",
"apac.amazon.nova-lite-v1:0",
"apac.amazon.nova-pro-v1:0",
"bedrock/us-gov-west-1/amazon.nova-micro-v1:0",
"bedrock/us-gov-west-1/amazon.nova-lite-v1:0",
"bedrock/us-gov-west-1/amazon.nova-pro-v1:0",
"bedrock/us-gov-east-1/amazon.nova-pro-v1:0",
],
)
def test_nova_prompt_caching_models_price_cache_reads_at_a_quarter_of_input(model, monkeypatch):
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
entry = litellm.model_cost[model]
assert entry["supports_prompt_caching"] is True
assert entry["cache_read_input_token_cost"] == pytest.approx(entry["input_cost_per_token"] * 0.25)
def test_transform_usage_with_reasoning_content():
"""Test that completion_tokens_details correctly tracks reasoning vs text tokens."""
usage = ConverseTokenUsageBlock(

View file

@ -324,18 +324,18 @@ CONVERSE_METADATA_EVENT = {
}
def _converse_stream_wrapper(events):
def _converse_stream_wrapper(events, model=CONVERSE_MODEL):
async def bedrock_stream():
decoder = AWSEventStreamDecoder(model=CONVERSE_MODEL)
decoder = AWSEventStreamDecoder(model=model)
for event in events:
yield decoder._chunk_parser(chunk_data=event)
return CustomStreamWrapper(
completion_stream=bedrock_stream(),
model=CONVERSE_MODEL,
model=model,
custom_llm_provider="bedrock",
logging_obj=LiteLLMLoggingObj(
model=CONVERSE_MODEL,
model=model,
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion",
@ -427,6 +427,46 @@ async def test_converse_stream_ends_on_finish_reason_chunk(events, expected_fini
assert any(getattr(chunk, "usage", None) is not None for chunk in wrapper.chunks)
@pytest.mark.asyncio
async def test_nova_invoke_stream_reports_bedrock_usage_and_finish_reason():
"""InvokeModel Nova wraps every Converse event under its event-type key and reports usage
without ``totalTokens``; the stream must end on Bedrock's finish reason and surface the
cached tokens instead of a token-count estimate."""
events = (
{"messageStart": {"role": "assistant"}},
{"contentBlockDelta": {"delta": {"text": "OK"}, "contentBlockIndex": 0}},
{"contentBlockDelta": {"delta": {"text": "."}, "contentBlockIndex": 0}},
{"contentBlockStop": {"contentBlockIndex": 0}},
{"messageStop": {"stopReason": "end_turn"}},
{
"metadata": {
"usage": {
"inputTokens": 5,
"outputTokens": 3,
"cacheReadInputTokenCount": 12262,
"cacheWriteInputTokenCount": 0,
},
"metrics": {},
"trace": {},
}
},
)
wrapper = _converse_stream_wrapper(events, model="bedrock/invoke/us.amazon.nova-pro-v1:0")
chunks = [chunk async for chunk in wrapper]
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "OK."
finish_reasons = [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason]
assert finish_reasons == ["stop"]
assert chunks[-1].choices[0].finish_reason == "stop"
usages = [chunk.usage for chunk in wrapper.chunks if getattr(chunk, "usage", None) is not None]
assert len(usages) == 1
assert usages[0].prompt_tokens == 12267
assert usages[0].prompt_tokens_details.cached_tokens == 12262
assert usages[0].completion_tokens == 3
assert usages[0].total_tokens == 12270
@pytest.mark.asyncio
async def test_converse_stream_still_emits_guardrail_trace_after_finish_reason():
"""Guardrail metadata events carry a trace payload alongside usage; that chunk must still reach the caller