diff --git a/tests/litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/litellm/llms/bedrock/chat/test_converse_transformation.py new file mode 100644 index 0000000000..7218063682 --- /dev/null +++ b/tests/litellm/llms/bedrock/chat/test_converse_transformation.py @@ -0,0 +1,36 @@ +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path +from unittest.mock import MagicMock, patch + +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.types.llms.bedrock import ConverseTokenUsageBlock + + +def test_transform_usage(): + usage = ConverseTokenUsageBlock( + **{ + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 10, + "cacheCreationInputTokenCount": 0, + "cacheCreationInputTokens": 0, + "inputTokens": 12, + "outputTokens": 56, + "totalTokens": 78, + } + ) + 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 diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index ca9e47f45a..e2948789fc 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2948,12 +2948,3 @@ async def test_bedrock_stream_thinking_content_openwebui(): assert ( len(response_content) > 0 ), "There should be non-empty content after thinking tags" - - -def test_bedrock_usage_block(): - litellm._turn_on_debug() - response = completion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", - messages=[{"role": "user", "content": "Hello who is this?"}], - ) - assert response.usage.total_tokens > 0