mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge pull request #41343 from BerriAI/litellm_lit5030_bedrock_invoke_nova_prompt_caching
fix(bedrock): make prompt caching work on the Nova InvokeModel route
This commit is contained in:
Mateo Wang
GitHub
commit
8e524370e1
9 changed files with 397 additions and 35 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,50 @@ 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_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,7 +99,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 +107,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 +147,24 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
|
|||
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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -2905,6 +2908,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,
|
||||
|
|
@ -2920,6 +2924,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,
|
||||
|
|
@ -2933,6 +2938,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,
|
||||
|
|
@ -12864,6 +12870,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,
|
||||
|
|
@ -13044,6 +13051,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,
|
||||
|
|
@ -13059,6 +13067,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,
|
||||
|
|
@ -13072,6 +13081,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,
|
||||
|
|
@ -21612,6 +21622,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,
|
||||
|
|
@ -21627,6 +21638,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,
|
||||
|
|
@ -21640,6 +21652,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,
|
||||
|
|
@ -45146,6 +45159,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,
|
||||
|
|
@ -45161,6 +45175,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,
|
||||
|
|
@ -45189,6 +45204,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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -2905,6 +2908,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,
|
||||
|
|
@ -2920,6 +2924,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,
|
||||
|
|
@ -2933,6 +2938,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,
|
||||
|
|
@ -12864,6 +12870,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,
|
||||
|
|
@ -13044,6 +13051,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,
|
||||
|
|
@ -13059,6 +13067,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,
|
||||
|
|
@ -13072,6 +13081,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,
|
||||
|
|
@ -21612,6 +21622,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,
|
||||
|
|
@ -21627,6 +21638,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,
|
||||
|
|
@ -21640,6 +21652,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,
|
||||
|
|
@ -45146,6 +45159,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,
|
||||
|
|
@ -45161,6 +45175,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,
|
||||
|
|
@ -45189,6 +45204,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,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,95 @@
|
|||
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": {}}}}]
|
||||
TOOL_CALL = {"id": "call_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}
|
||||
PNG_DATA_URL = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
|
||||
|
||||
|
||||
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_behind_a_non_text_block_moves_back_to_the_last_text_block(local_model_cost_map):
|
||||
"""InvokeModel rejects ``cachePoint`` on image, toolUse, and toolResult blocks
|
||||
(``extraneous key [cachePoint] is not permitted``), so the point a user put on an image or a
|
||||
tool result lands on the closest text block before it, and a message with no text block at
|
||||
all sends no point rather than a request AWS refuses.
|
||||
"""
|
||||
request = _transform_request(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "what is in this picture?"},
|
||||
{"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "cache_control": EPHEMERAL},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": None, "tool_calls": [TOOL_CALL]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "sunny", "cache_control": EPHEMERAL},
|
||||
],
|
||||
optional_params={"tools": TOOLS},
|
||||
)
|
||||
picture, image = request["messages"][0]["content"]
|
||||
assert picture == {"text": "what is in this picture?", "cachePoint": DEFAULT_CACHE_POINT}
|
||||
assert set(image) == {"image"}
|
||||
assert [set(block) for block in request["messages"][2]["content"]] == [{"toolResult"}]
|
||||
|
||||
|
||||
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
|
||||
|
|
@ -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 the entry's discounted cache read 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 0 < model_info["cache_read_input_token_cost"] < model_info["input_cost_per_token"]
|
||||
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_below_the_input_rate(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 0 < entry["cache_read_input_token_cost"] < entry["input_cost_per_token"]
|
||||
|
||||
|
||||
def test_transform_usage_with_reasoning_content():
|
||||
"""Test that completion_tokens_details correctly tracks reasoning vs text tokens."""
|
||||
usage = ConverseTokenUsageBlock(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue