misc/prompt_caching.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prompt caching through the Anthropic API\n",
    "\n",
    "Prompt caching allows you to store and reuse context within your prompt. This makes it more practical to include additional information in your prompt—such as detailed instructions and example responses—which help improve every response Claude generates.\n",
    "\n",
    "In addition, by fully leveraging prompt caching within your prompt, you can reduce latency by >2x and costs up to 90%. This can generate significant savings when building solutions that involve repetitive tasks around detailed book_content.\n",
    "\n",
    "In this cookbook, we will demonstrate how to use prompt caching in a single turn and across a multi-turn conversation. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "## Setup\n",
    "\n",
    "First, let's set up our environment with the necessary imports and initializations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install anthropic bs4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import anthropic\n",
    "import time\n",
    "import requests\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "client = anthropic.Anthropic()\n",
    "MODEL_NAME = \"claude-3-5-sonnet-20240620\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's fetch some text content to use in our examples. We'll use the text from Pride and Prejudice by Jane Austen which is around ~187,000 tokens long."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fetched 737522 characters from the book.\n",
      "First 500 characters:\n",
      "The Project Gutenberg eBook of Pride and Prejudice\n",
      "This ebook is for the use of anyone anywhere in the United States and\n",
      "most other parts of the world at no cost and with almost no restrictions\n",
      "whatsoever. You may copy it, give it away or re-use it under the terms\n",
      "of the Project Gutenberg License included with this ebook or online\n",
      "at www.gutenberg.org. If you are not located in the United States,\n",
      "you will have to check the laws of the country where you are located\n",
      "before using this eBook.\n",
      "Title:\n"
     ]
    }
   ],
   "source": [
    "def fetch_article_content(url):\n",
    "    response = requests.get(url)\n",
    "    soup = BeautifulSoup(response.content, 'html.parser')\n",
    "    \n",
    "    # Remove script and style elements\n",
    "    for script in soup([\"script\", \"style\"]):\n",
    "        script.decompose()\n",
    "    \n",
    "    # Get text\n",
    "    text = soup.get_text()\n",
    "    \n",
    "    # Break into lines and remove leading and trailing space on each\n",
    "    lines = (line.strip() for line in text.splitlines())\n",
    "    # Break multi-headlines into a line each\n",
    "    chunks = (phrase.strip() for line in lines for phrase in line.split(\"  \"))\n",
    "    # Drop blank lines\n",
    "    text = '\\n'.join(chunk for chunk in chunks if chunk)\n",
    "    \n",
    "    return text\n",
    "\n",
    "# Fetch the content of the article\n",
    "book_url = \"https://www.gutenberg.org/cache/epub/1342/pg1342.txt\"\n",
    "book_content = fetch_article_content(book_url)\n",
    "\n",
    "print(f\"Fetched {len(book_content)} characters from the book.\")\n",
    "print(\"First 500 characters:\")\n",
    "print(book_content[:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Single turn\n",
    "\n",
    "Let's demonstrate prompt caching with a large document, comparing the performance and cost between cached and non-cached API calls."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 1: Non-cached API Call\n",
    "\n",
    "First, let's make a non-cached API call. This will load the prompt into the cache so that our subsequent cached API calls can benefit from the prompt caching.\n",
    "\n",
    "We will ask for a short output string to keep the output response time low since the benefit of prompt caching applies only to the input processing time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Non-cached API call time: 21.44 seconds\n",
      "Non-cached API call input tokens: 17\n",
      "Non-cached API call output tokens: 8\n",
      "\n",
      "Summary (non-cached):\n",
      "[TextBlock(text='Pride and Prejudice', type='text')]\n"
     ]
    }
   ],
   "source": [
    "def make_non_cached_api_call():\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": \"<book>\" + book_content + \"</book>\",\n",
    "                    \"cache_control\": {\"type\": \"ephemeral\"}\n",
    "                },\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": \"What is the title of this book? Only output the title.\"\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    start_time = time.time()\n",
    "    response = client.messages.create(\n",
    "        model=MODEL_NAME,\n",
    "        max_tokens=300,\n",
    "        messages=messages,\n",
    "        extra_headers={\"anthropic-beta\": \"prompt-caching-2024-07-31\"}\n",
    "\n",
    "    )\n",
    "    end_time = time.time()\n",
    "\n",
    "    return response, end_time - start_time\n",
    "\n",
    "non_cached_response, non_cached_time = make_non_cached_api_call()\n",
    "\n",
    "print(f\"Non-cached API call time: {non_cached_time:.2f} seconds\")\n",
    "print(f\"Non-cached API call input tokens: {non_cached_response.usage.input_tokens}\")\n",
    "print(f\"Non-cached API call output tokens: {non_cached_response.usage.output_tokens}\")\n",
    "\n",
    "print(\"\\nSummary (non-cached):\")\n",
    "print(non_cached_response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Part 2: Cached API Call\n",
    "\n",
    "Now, let's make a cached API call. I'll add in the \"cache_control\": {\"type\": \"ephemeral\"} attribute to the content object and add the \"prompt-caching-2024-07-31\" beta header to the request. This will enable prompt caching for this API call.\n",
    "\n",
    "To keep the output latency constant, we will ask Claude the same question as before. Note that this question is not part of the cached content."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cached API call time: 3.64 seconds\n",
      "Cached API call input tokens: 17\n",
      "Cached API call output tokens: 16\n",
      "\n",
      "Summary (cached):\n",
      "[TextBlock(text='The title of this book is \"Pride and Prejudice\".', type='text')]\n"
     ]
    }
   ],
   "source": [
    "def make_cached_api_call():\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": \"<book>\" + book_content + \"</book>\",\n",
    "                    \"cache_control\": {\"type\": \"ephemeral\"}\n",
    "                },\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": \"What is the title of this book? Only output the title.\"\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    start_time = time.time()\n",
    "    response = client.messages.create(\n",
    "        model=MODEL_NAME,\n",
    "        max_tokens=300,\n",
    "        messages=messages,\n",
    "        extra_headers={\"anthropic-beta\": \"prompt-caching-2024-07-31\"}\n",
    "    )\n",
    "    end_time = time.time()\n",
    "\n",
    "    return response, end_time - start_time\n",
    "\n",
    "cached_response, cached_time = make_cached_api_call()\n",
    "\n",
    "print(f\"Cached API call time: {cached_time:.2f} seconds\")\n",
    "print(f\"Cached API call input tokens: {cached_response.usage.input_tokens}\")\n",
    "print(f\"Cached API call output tokens: {cached_response.usage.output_tokens}\")\n",
    "\n",
    "print(\"\\nSummary (cached):\")\n",
    "print(cached_response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the cached API call only took 3.64 seconds total compared to 21.44 seconds for the non-cached API call. This is a significant improvement in overall latency due to caching."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Multi-turn Conversation with Incremental Caching\n",
    "\n",
    "Now, let's look at a multi-turn conversation where we add cache breakpoints as the conversation progresses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Turn 1:\n",
      "User: What is the title of this novel?\n",
      "Assistant: The title of this novel is \"Pride and Prejudice\" by Jane Austen.\n",
      "User input tokens: 4\n",
      "Output tokens: 22\n",
      "Input tokens (cache read): 0\n",
      "Input tokens (cache write): 187354\n",
      "0.0% of input prompt cached (4 tokens)\n",
      "Time taken: 23.98 seconds\n",
      "\n",
      "Turn 2:\n",
      "User: Who are Mr. and Mrs. Bennet?\n",
      "Assistant: Mr. and Mrs. Bennet are the parents of the five Bennet daughters in the novel. They are key characters in the story:\n",
      "\n",
      "1. Mr. Bennet is the father of the family. He is described as an intelligent but sarcastic man who is often amused by the folly of others, including his wife. He tends to be detached from family matters and retreats to his library for peace.\n",
      "\n",
      "2. Mrs. Bennet is the mother. She is described as a woman \"of mean understanding, little information, and uncertain temper.\" Her main goal in life is to see her daughters married to wealthy men. She is often portrayed as silly, nervous, and prone to hyperbole.\n",
      "\n",
      "The Bennets have five daughters: Jane, Elizabeth (the protagonist), Mary, Kitty, and Lydia. The family lives at Longbourn estate, which is entailed to Mr. Collins (a male cousin) because Mr. and Mrs. Bennet have no sons.\n",
      "User input tokens: 4\n",
      "Output tokens: 218\n",
      "Input tokens (cache read): 187354\n",
      "Input tokens (cache write): 36\n",
      "100.0% of input prompt cached (187358 tokens)\n",
      "Time taken: 8.11 seconds\n",
      "\n",
      "Turn 3:\n",
      "User: What is Netherfield Park?\n",
      "Assistant: Netherfield Park is an estate near the Bennet family's home of Longbourn in the novel. It plays a significant role in the story for several reasons:\n",
      "\n",
      "1. It is rented by Mr. Bingley, a wealthy young man who becomes a central character in the novel.\n",
      "\n",
      "2. The arrival of Mr. Bingley at Netherfield sets much of the plot in motion, as he becomes a potential suitor for Jane Bennet.\n",
      "\n",
      "3. It's where several important events take place, including a ball where Elizabeth and Mr. Darcy have significant interactions.\n",
      "\n",
      "4. Jane Bennet stays at Netherfield for a period when she falls ill, allowing for further development of relationships between the characters.\n",
      "\n",
      "5. It serves as a symbol of wealth and social status in the story, contrasting with the more modest Longbourn estate of the Bennets.\n",
      "\n",
      "The comings and goings at Netherfield Park, and the social events held there, are crucial to the development of the romantic relationships and social dynamics in the novel.\n",
      "User input tokens: 4\n",
      "Output tokens: 234\n",
      "Input tokens (cache read): 187390\n",
      "Input tokens (cache write): 229\n",
      "100.0% of input prompt cached (187394 tokens)\n",
      "Time taken: 8.24 seconds\n",
      "\n",
      "Turn 4:\n",
      "User: What is the main theme of this novel?\n",
      "Assistant: The main theme of \"Pride and Prejudice\" is reflected in its title: the interplay between pride and prejudice in relationships and society. However, there are several important themes that run throughout the novel:\n",
      "\n",
      "1. Love and Marriage: The novel explores various attitudes towards marriage, from practical arrangements to romantic love.\n",
      "\n",
      "2. Pride and Prejudice: How first impressions and preconceived notions can lead to misunderstandings and misjudgments.\n",
      "\n",
      "3. Class and Social Status: The impact of social hierarchy on relationships and personal choices.\n",
      "\n",
      "4. Personal Growth and Self-awareness: Particularly through the character development of Elizabeth and Darcy.\n",
      "\n",
      "5. The Role of Women in Society: The novel examines the limited options available to women in early 19th-century England.\n",
      "\n",
      "6. Family Relationships: The dynamics within the Bennet family and how they influence the characters' decisions.\n",
      "\n",
      "7. Reputation and Honor: How these concepts affect social standing and personal relationships.\n",
      "\n",
      "8. The Balance between Sense and Sensibility: Finding a middle ground between emotional impulses and rational thought.\n",
      "\n",
      "While all these themes are present, the central focus is on how pride and prejudice influence relationships, particularly in the developing romance between Elizabeth Bennet and Mr. Darcy. Their journey from mutual dislike to understanding and love forms the core of the novel's thematic exploration.\n",
      "User input tokens: 4\n",
      "Output tokens: 299\n",
      "Input tokens (cache read): 187619\n",
      "Input tokens (cache write): 247\n",
      "100.0% of input prompt cached (187623 tokens)\n",
      "Time taken: 9.75 seconds\n"
     ]
    }
   ],
   "source": [
    "class ConversationHistory:\n",
    "    def __init__(self):\n",
    "        # Initialize an empty list to store conversation turns\n",
    "        self.turns = []\n",
    "\n",
    "    def add_turn_assistant(self, content):\n",
    "        # Add an assistant's turn to the conversation history\n",
    "        self.turns.append({\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": content\n",
    "                }\n",
    "            ]\n",
    "        })\n",
    "\n",
    "    def add_turn_user(self, content):\n",
    "        # Add a user's turn to the conversation history\n",
    "        self.turns.append({\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": content\n",
    "                }\n",
    "            ]\n",
    "        })\n",
    "\n",
    "    def get_turns(self):\n",
    "        # Retrieve conversation turns with specific formatting\n",
    "        result = []\n",
    "        user_turns_processed = 0\n",
    "        # Iterate through turns in reverse order\n",
    "        for turn in reversed(self.turns):\n",
    "            if turn[\"role\"] == \"user\" and user_turns_processed < 2:\n",
    "                # Add the last two user turns with ephemeral cache control\n",
    "                result.append({\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": [\n",
    "                        {\n",
    "                            \"type\": \"text\",\n",
    "                            \"text\": turn[\"content\"][0][\"text\"],\n",
    "                            \"cache_control\": {\"type\": \"ephemeral\"}\n",
    "                        }\n",
    "                    ]\n",
    "                })\n",
    "                user_turns_processed += 1\n",
    "            else:\n",
    "                # Add other turns as they are\n",
    "                result.append(turn)\n",
    "        # Return the turns in the original order\n",
    "        return list(reversed(result))\n",
    "\n",
    "# Initialize the conversation history\n",
    "conversation_history = ConversationHistory()\n",
    "\n",
    "# System message containing the book content\n",
    "# Note: 'book_content' should be defined elsewhere in the code\n",
    "system_message = f\"<file_contents> {book_content} </file_contents>\"\n",
    "\n",
    "# Predefined questions for our simulation\n",
    "questions = [\n",
    "    \"What is the title of this novel?\",\n",
    "    \"Who are Mr. and Mrs. Bennet?\",\n",
    "    \"What is Netherfield Park?\",\n",
    "    \"What is the main theme of this novel?\"\n",
    "]\n",
    "\n",
    "def simulate_conversation():\n",
    "    for i, question in enumerate(questions, 1):\n",
    "        print(f\"\\nTurn {i}:\")\n",
    "        print(f\"User: {question}\")\n",
    "        \n",
    "        # Add user input to conversation history\n",
    "        conversation_history.add_turn_user(question)\n",
    "\n",
    "        # Record the start time for performance measurement\n",
    "        start_time = time.time()\n",
    "\n",
    "        # Make an API call to the assistant\n",
    "        response = client.messages.create(\n",
    "            model=MODEL_NAME,\n",
    "            extra_headers={\n",
    "              \"anthropic-beta\": \"prompt-caching-2024-07-31\"\n",
    "            },\n",
    "            max_tokens=300,\n",
    "            system=[\n",
    "                {\"type\": \"text\", \"text\": system_message, \"cache_control\": {\"type\": \"ephemeral\"}},\n",
    "            ],\n",
    "            messages=conversation_history.get_turns(),\n",
    "        )\n",
    "\n",
    "        # Record the end time\n",
    "        end_time = time.time()\n",
    "\n",
    "        # Extract the assistant's reply\n",
    "        assistant_reply = response.content[0].text\n",
    "        print(f\"Assistant: {assistant_reply}\")\n",
    "\n",
    "        # Print token usage information\n",
    "        input_tokens = response.usage.input_tokens\n",
    "        output_tokens = response.usage.output_tokens\n",
    "        input_tokens_cache_read = getattr(response.usage, 'cache_read_input_tokens', '---')\n",
    "        input_tokens_cache_create = getattr(response.usage, 'cache_creation_input_tokens', '---')\n",
    "        print(f\"User input tokens: {input_tokens}\")\n",
    "        print(f\"Output tokens: {output_tokens}\")\n",
    "        print(f\"Input tokens (cache read): {input_tokens_cache_read}\")\n",
    "        print(f\"Input tokens (cache write): {input_tokens_cache_create}\")\n",
    "\n",
    "        # Calculate and print the elapsed time\n",
    "        elapsed_time = end_time - start_time\n",
    "\n",
    "        # Calculate the percentage of input prompt cached\n",
    "        total_input_tokens = input_tokens + (int(input_tokens_cache_read) if input_tokens_cache_read != '---' else 0)\n",
    "        percentage_cached = (int(input_tokens_cache_read) / total_input_tokens * 100 if input_tokens_cache_read != '---' and total_input_tokens > 0 else 0)\n",
    "\n",
    "        print(f\"{percentage_cached:.1f}% of input prompt cached ({total_input_tokens} tokens)\")\n",
    "        print(f\"Time taken: {elapsed_time:.2f} seconds\")\n",
    "\n",
    "        # Add assistant's reply to conversation history\n",
    "        conversation_history.add_turn_assistant(assistant_reply)\n",
    "\n",
    "# Run the simulated conversation\n",
    "simulate_conversation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see in this example, response times decreased from nearly 24 seconds to just 8-9 seconds after the initial cache setup, while maintaining the same level of quality across the answers. Most of this remaining latency is due to the time it takes to generate the response, which is not affected by prompt caching.\n",
    "\n",
    "And since nearly 100% of input tokens were cached in subsequent turns as we kept adjusting the cache breakpoints, we were able to read the next user message nearly instantly."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}