diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 540ec13f80..bb874cfe38 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -612,8 +612,8 @@ class AmazonConverseConfig(BaseConfig): if "cacheReadInputTokens" in usage: cache_read_input_tokens = usage["cacheReadInputTokens"] input_tokens += cache_read_input_tokens - if "cacheCreationInputTokens" in usage: - cache_creation_input_tokens = usage["cacheCreationInputTokens"] + if "cacheWriteInputTokens" in usage: + cache_creation_input_tokens = usage["cacheWriteInputTokens"] input_tokens += cache_creation_input_tokens prompt_tokens_details = PromptTokensDetailsWrapper( diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 836efc44ee..57fb04c8a9 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -111,8 +111,8 @@ class ConverseTokenUsageBlock(TypedDict): totalTokens: int cacheReadInputTokenCount: int cacheReadInputTokens: int - cacheCreationInputTokenCount: int - cacheCreationInputTokens: int + cacheWriteInputTokenCount: int + cacheWriteInputTokens: int class ConverseResponseBlock(TypedDict): diff --git a/tests/litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/litellm/llms/bedrock/chat/test_converse_transformation.py index 7218063682..e912ada8ff 100644 --- a/tests/litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/litellm/llms/bedrock/chat/test_converse_transformation.py @@ -18,19 +18,27 @@ def test_transform_usage(): usage = ConverseTokenUsageBlock( **{ "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 10, - "cacheCreationInputTokenCount": 0, - "cacheCreationInputTokens": 0, - "inputTokens": 12, - "outputTokens": 56, - "totalTokens": 78, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 1789, + "cacheWriteInputTokens": 1789, + "inputTokens": 3, + "outputTokens": 401, + "totalTokens": 2193, } ) config = AmazonConverseConfig() openai_usage = config._transform_usage(usage) - assert openai_usage.prompt_tokens == 22 - assert openai_usage.completion_tokens == 56 - assert openai_usage.total_tokens == 78 - assert openai_usage.prompt_tokens_details.cached_tokens == 10 - assert openai_usage._cache_creation_input_tokens == 0 - assert openai_usage._cache_read_input_tokens == 10 + assert ( + openai_usage.prompt_tokens + == usage["inputTokens"] + + usage["cacheWriteInputTokens"] + + usage["cacheReadInputTokens"] + ) + assert openai_usage.completion_tokens == usage["outputTokens"] + assert openai_usage.total_tokens == usage["totalTokens"] + assert ( + openai_usage.prompt_tokens_details.cached_tokens + == usage["cacheReadInputTokens"] + ) + assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"] + assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"]