From b0651bd481fa32658364aee65a40baa56d51dc1f Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 14 Aug 2024 14:56:49 -0700
Subject: [PATCH] add anthropic cache controls

---
 litellm/llms/prompt_templates/factory.py | 62 +++++++++++++++++++-----
 litellm/types/llms/anthropic.py          |  6 ++-
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py
index 7c3c7e80fb..66658e23a4 100644
--- a/litellm/llms/prompt_templates/factory.py
+++ b/litellm/llms/prompt_templates/factory.py
@@ -1224,6 +1224,19 @@ def convert_to_anthropic_tool_invoke(
     return anthropic_tool_invoke
 
 
+def add_cache_control_to_content(
+    anthropic_content_element: Union[
+        dict, AnthropicMessagesImageParam, AnthropicMessagesTextParam
+    ],
+    orignal_content_element: dict,
+):
+    if "cache_control" in orignal_content_element:
+        anthropic_content_element["cache_control"] = orignal_content_element[
+            "cache_control"
+        ]
+    return anthropic_content_element
+
+
 def anthropic_messages_pt(
     messages: list,
     model: str,
@@ -1264,18 +1277,31 @@ def anthropic_messages_pt(
                         image_chunk = convert_to_anthropic_image_obj(
                             m["image_url"]["url"]
                         )
-                        user_content.append(
-                            AnthropicMessagesImageParam(
-                                type="image",
-                                source=AnthropicImageParamSource(
-                                    type="base64",
-                                    media_type=image_chunk["media_type"],
-                                    data=image_chunk["data"],
-                                ),
-                            )
+
+                        _anthropic_content_element = AnthropicMessagesImageParam(
+                            type="image",
+                            source=AnthropicImageParamSource(
+                                type="base64",
+                                media_type=image_chunk["media_type"],
+                                data=image_chunk["data"],
+                            ),
                         )
+
+                        anthropic_content_element = add_cache_control_to_content(
+                            anthropic_content_element=_anthropic_content_element,
+                            orignal_content_element=m,
+                        )
+                        user_content.append(anthropic_content_element)
                     elif m.get("type", "") == "text":
-                        user_content.append({"type": "text", "text": m["text"]})
+                        _anthropic_text_content_element = {
+                            "type": "text",
+                            "text": m["text"],
+                        }
+                        anthropic_content_element = add_cache_control_to_content(
+                            anthropic_content_element=_anthropic_text_content_element,
+                            orignal_content_element=m,
+                        )
+                        user_content.append(anthropic_content_element)
             elif (
                 messages[msg_i]["role"] == "tool"
                 or messages[msg_i]["role"] == "function"
@@ -1306,6 +1332,10 @@ def anthropic_messages_pt(
                         anthropic_message = AnthropicMessagesTextParam(
                             type="text", text=m.get("text")
                         )
+                        anthropic_message = add_cache_control_to_content(
+                            anthropic_content_element=anthropic_message,
+                            orignal_content_element=m,
+                        )
                         assistant_content.append(anthropic_message)
             elif (
                 "content" in messages[msg_i]
@@ -1313,9 +1343,17 @@ def anthropic_messages_pt(
                 and len(messages[msg_i]["content"])
                 > 0  # don't pass empty text blocks. anthropic api raises errors.
             ):
-                assistant_content.append(
-                    {"type": "text", "text": messages[msg_i]["content"]}
+
+                _anthropic_text_content_element = {
+                    "type": "text",
+                    "text": messages[msg_i]["content"],
+                }
+
+                anthropic_content_element = add_cache_control_to_content(
+                    anthropic_content_element=_anthropic_text_content_element,
+                    orignal_content_element=messages[msg_i],
                 )
+                assistant_content.append(anthropic_content_element)
 
             if messages[msg_i].get(
                 "tool_calls", []
diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py
index 36bcb6cc73..2eb2aef549 100644
--- a/litellm/types/llms/anthropic.py
+++ b/litellm/types/llms/anthropic.py
@@ -15,9 +15,10 @@ class AnthropicMessagesTool(TypedDict, total=False):
     input_schema: Required[dict]
 
 
-class AnthropicMessagesTextParam(TypedDict):
+class AnthropicMessagesTextParam(TypedDict, total=False):
     type: Literal["text"]
     text: str
+    cache_control: Optional[dict]
 
 
 class AnthropicMessagesToolUseParam(TypedDict):
@@ -54,9 +55,10 @@ class AnthropicImageParamSource(TypedDict):
     data: str
 
 
-class AnthropicMessagesImageParam(TypedDict):
+class AnthropicMessagesImageParam(TypedDict, total=False):
     type: Literal["image"]
     source: AnthropicImageParamSource
+    cache_control: Optional[dict]
 
 
 class AnthropicMessagesToolResultContent(TypedDict):