From d43441ae5d7fece4ca7ac97e919398c3cd195d1e Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 27 Aug 2024 13:57:03 -0700
Subject: [PATCH] fix(anthropic.py): support setting cache control headers,
 automatically

Don't require user to manually pass in 'extra_headers' for anthropic cache control usage
---
 litellm/llms/anthropic.py                     | 33 +++++++++-
 .../tests/test_anthropic_prompt_caching.py    | 60 +++++++++++++++++++
 2 files changed, 91 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/anthropic.py b/litellm/llms/anthropic.py
index 9f1346b0c33..813897c663c 100644
--- a/litellm/llms/anthropic.py
+++ b/litellm/llms/anthropic.py
@@ -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)
diff --git a/litellm/tests/test_anthropic_prompt_caching.py b/litellm/tests/test_anthropic_prompt_caching.py
index 87bfc23f841..4b7a4a3ded5 100644
--- a/litellm/tests/test_anthropic_prompt_caching.py
+++ b/litellm/tests/test_anthropic_prompt_caching.py
@@ -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