scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

test test_anthropic_api_prompt_caching_basic

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-14 15:06:10 -07:00
parent 69a640e9c4
commit 96f9655a02
2 changed files with 15 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/llms/anthropic.py
View file

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

10
litellm/tests/test_completion.py
View file

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