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

litellm/llms/bedrock/chat/converse_transformation.py

@@ -1902,7 +1902,7 @@ def _parse_cache_details(usage: ConverseTokenUsageBlock) -> "CacheCreationTokenD
             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 @@ def _usage_count(usage_object: Mapping[str, object], *keys: str) -> int:
                 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 @@ def usage_from_batch_output(self, usage_object: Mapping[str, object]) -> Usage:
         """
         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 @@ def transform_usage(
         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,

litellm/llms/bedrock/chat/invoke_handler.py

@@ -3,6 +3,7 @@
 from typing import Final, cast
 
 import httpx
+from pydantic import TypeAdapter
 
 import litellm
 from litellm import verbose_logger
@@ -51,6 +52,15 @@
 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 @@ def converse_chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
             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 @@ def _chunk_parser(self, chunk_data: dict) -> GChunk | ModelResponseStream | dict
         ):
             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:

litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py

@@ -6,12 +6,21 @@
 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,50 @@
 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_block_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)
+        anchor: Final = next((index for index in reversed(range(len(inlined))) if "text" in inlined[index]), None)
+        if anchor is None:
+            return inlined
+        return (*inlined[:anchor], with_cache_point(inlined[anchor], cache_point), *inlined[anchor + 1 :])
+
+    return list(reduce(attach, blocks, ()))
+
 
 class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
     """
@@ -46,19 +99,21 @@ def transform_request(
         self,
         model: str,
         messages: list[AllMessageValues],
-        optional_params: dict,
+        optional_params: dict[str, object],
         litellm_params: dict,
         headers: dict,
     ) -> dict:
         _transformed_nova_request: Final = AmazonConverseConfig.transform_request(
             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 +147,24 @@ def transform_response(
             json_mode,
         )
 
+    @staticmethod
+    def _inline_cache_points(request: BedrockInvokeNovaRequest) -> BedrockInvokeNovaRequest:
+        """InvokeModel takes ``cachePoint`` as a key of the text block it caches: it rejects the
+        standalone ``{"cachePoint": ...}`` blocks Converse accepts and the key on image, toolUse,
+        and toolResult blocks, so a point behind one of those moves back to the last text block.
+        """
+        return {
+            **request,
+            "system": _inline_block_cache_points(request.get("system", []), _system_block_with_cache_point),
+            "messages": [
+                MessageBlock(
+                    role=message["role"],
+                    content=_inline_block_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.

litellm/model_prices_and_context_window_backup.json

@@ -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,

litellm/types/llms/bedrock.py

@@ -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]