fix(anthropic.py): support setting cache control headers, automatically

Don't require user to manually pass in 'extra_headers' for anthropic cache control usage
This commit is contained in:
Krrish Dholakia 2024-08-27 13:57:03 -07:00
parent 63adb3f940
commit d43441ae5d
2 changed files with 91 additions and 2 deletions

View file

@ -158,6 +158,12 @@ class AnthropicConfig:
"extra_headers",
]
def get_cache_control_headers(self) -> dict:
return {
"anthropic-version": "2023-06-01",
"anthropic-beta": "prompt-caching-2024-07-31",
}
def map_openai_params(self, non_default_params: dict, optional_params: dict):
for param, value in non_default_params.items():
if param == "max_tokens":
@ -203,6 +209,20 @@ class AnthropicConfig:
optional_params["top_p"] = value
return optional_params
def is_cache_control_set(self, messages: List[AllMessageValues]) -> bool:
"""
Return if {"cache_control": ..} in message content block
Used to check if anthropic prompt caching headers need to be set.
"""
for message in messages:
if message["content"] is not None and isinstance(message["content"], list):
for content in message["content"]:
if "cache_control" in content:
return True
return False
### FOR [BETA] `/v1/messages` endpoint support
def translatable_anthropic_params(self) -> List:
@ -563,19 +583,28 @@ class AnthropicConfig:
# makes headers for API call
def validate_environment(api_key, user_headers, model):
def validate_environment(
api_key, user_headers, model, messages: List[AllMessageValues]
):
cache_headers = {}
if api_key is None:
raise litellm.AuthenticationError(
message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` in your environment vars",
llm_provider="anthropic",
model=model,
)
if AnthropicConfig().is_cache_control_set(messages=messages):
cache_headers = AnthropicConfig().get_cache_control_headers()
headers = {
"accept": "application/json",
"anthropic-version": "2023-06-01",
"content-type": "application/json",
"x-api-key": api_key,
}
headers.update(cache_headers)
if user_headers is not None and isinstance(user_headers, dict):
headers = {**headers, **user_headers}
return headers
@ -891,7 +920,7 @@ class AnthropicChatCompletion(BaseLLM):
headers={},
client=None,
):
headers = validate_environment(api_key, headers, model)
headers = validate_environment(api_key, headers, model, messages=messages)
_is_function_call = False
messages = copy.deepcopy(messages)
optional_params = copy.deepcopy(optional_params)

View file

@ -222,6 +222,66 @@ async def test_anthropic_api_prompt_caching_basic():
)
@pytest.mark.asyncio()
async def test_anthropic_api_prompt_caching_no_headers():
litellm.set_verbose = True
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet-20240620",
messages=[
# System Message
{
"role": "system",
"content": [
{
"type": "text",
"text": "Here is the full text of a complex legal agreement"
* 400,
"cache_control": {"type": "ephemeral"},
}
],
},
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
{
"role": "assistant",
"content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
},
# The final turn is marked with cache-control, for continuing in followups.
{
"role": "user",
"content": [
{
"type": "text",
"text": "What are the key terms and conditions in this agreement?",
"cache_control": {"type": "ephemeral"},
}
],
},
],
temperature=0.2,
max_tokens=10,
)
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.asyncio
async def test_litellm_anthropic_prompt_caching_system():
# https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching#prompt-caching-examples