From 96f9655a029682d3f2ad5f893c0af38f3fe585c2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 14 Aug 2024 15:06:10 -0700 Subject: [PATCH] test test_anthropic_api_prompt_caching_basic --- litellm/llms/anthropic.py | 6 ++++++ litellm/tests/test_completion.py | 10 +++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py index fd4009b9736..c9f7856e9be 100644 --- a/litellm/llms/anthropic.py +++ b/litellm/llms/anthropic.py @@ -759,6 +759,7 @@ class AnthropicChatCompletion(BaseLLM): ## CALCULATING USAGE prompt_tokens = completion_response["usage"]["input_tokens"] completion_tokens = completion_response["usage"]["output_tokens"] + _usage = completion_response["usage"] total_tokens = prompt_tokens + completion_tokens model_response.created = int(time.time()) @@ -768,6 +769,11 @@ class AnthropicChatCompletion(BaseLLM): completion_tokens=completion_tokens, total_tokens=total_tokens, ) + + if "cache_creation_input_tokens" in _usage: + usage["cache_creation_input_tokens"] = _usage["cache_creation_input_tokens"] + if "cache_read_input_tokens" in _usage: + usage["cache_read_input_tokens"] = _usage["cache_read_input_tokens"] setattr(model_response, "usage", usage) # type: ignore return model_response diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 969805fb0a6..869339f786d 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3450,7 +3450,7 @@ def response_format_tests(response: litellm.ModelResponse): @pytest.mark.asyncio() -async def test_anthropic_api_prompt_caching_2(): +async def test_anthropic_api_prompt_caching_basic(): litellm.set_verbose = True response = await litellm.acompletion( model="anthropic/claude-3-5-sonnet-20240620", @@ -3504,6 +3504,14 @@ async def test_anthropic_api_prompt_caching_2(): print("response=", response) + assert "cache_read_input_tokens" in response.usage + assert "cache_creation_input_tokens" in response.usage + + # Assert either a cache entry was created or cache was read - changes depending on the anthropic api ttl + assert (response.usage.cache_read_input_tokens > 0) or ( + response.usage.cache_creation_input_tokens > 0 + ) + @pytest.mark.parametrize( "model",