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": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install anthropic bs4 --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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-20241022\""
   ]
  },
  {
   "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": 5,
   "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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Non-cached API call time: 20.37 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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cached API call time: 2.92 seconds\n",
      "Cached API call input tokens: 17\n",
      "Cached API call output tokens: 8\n",
      "\n",
      "Summary (cached):\n",
      "[TextBlock(text='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": 8,
   "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: 22.04 seconds\n",
      "\n",
      "Turn 2:\n",
      "User: Who are Mr. and Mrs. Bennet?\n",
      "Assistant: Mr. and Mrs. Bennet are the parents of five daughters (Jane, Elizabeth, Mary, Kitty, and Lydia) in Pride and Prejudice. Their personalities are quite different from each other:\n",
      "\n",
      "Mr. Bennet is an intelligent but rather detached and sarcastic man who often retreats to his library to avoid his wife's dramatic behavior. He is fond of his daughter Elizabeth (Lizzy) in particular due to her wit and good sense. While he can be amusing, he is also somewhat negligent in his duties as a father and head of the household.\n",
      "\n",
      "Mrs. Bennet is a rather foolish, nervous, and obsessive woman whose main goal in life is to see her daughters married to wealthy men. This preoccupation stems partly from the fact that their estate is entailed to a male heir (Mr. Collins), meaning her daughters will be left with little after Mr. Bennet's death. She often embarrasses her more sensible daughters (particularly Elizabeth and Jane) with her lack of social graces and her obvious attempts at matchmaking.\n",
      "\n",
      "Their marriage is portrayed as an unfortunate example of a match based on superficial attraction (Mr. Bennet was attracted to Mrs. Bennet's youth and beauty) rather than compatibility of character and intellect, serving as a contrast to the more meaningful relationships that develop later in the novel.\n",
      "User input tokens: 4\n",
      "Output tokens: 295\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.32 seconds\n",
      "\n",
      "Turn 3:\n",
      "User: What is Netherfield Park?\n",
      "Assistant: Netherfield Park is a large estate near the Bennets' home of Longbourn that is rented by Mr. Bingley at the beginning of the novel. It becomes a central location in the story as Mr. Bingley's arrival there sets many of the novel's key events in motion.\n",
      "\n",
      "The arrival of Mr. Bingley at Netherfield creates much excitement in the neighborhood, particularly for Mrs. Bennet, who sees him as a potential husband for one of her daughters due to his wealth (he has an income of around 4-5 thousand pounds a year). Indeed, it is at Netherfield where the romance between Jane Bennet and Mr. Bingley begins to develop.\n",
      "\n",
      "Netherfield is also significant because it is where Jane falls ill while visiting Mr. Bingley's sisters, leading to Elizabeth coming to nurse her. This extended stay at Netherfield provides important opportunities for Elizabeth to interact with Mr. Darcy, though at this point in the novel their interactions are marked by pride and prejudice on both sides.\n",
      "\n",
      "The house and estate represent wealth and social status in the novel, and its occupancy by the Bingleys brings the higher social circles of London into contact with the more provincial society of Meryton and its surroundings.\n",
      "User input tokens: 4\n",
      "Output tokens: 286\n",
      "Input tokens (cache read): 187390\n",
      "Input tokens (cache write): 306\n",
      "100.0% of input prompt cached (187394 tokens)\n",
      "Time taken: 10.73 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 the interplay between pride and prejudice in relationships, particularly as shown through the central romance between Elizabeth Bennet and Mr. Darcy. However, there are several important related themes:\n",
      "\n",
      "1. Pride and Prejudice:\n",
      "- Darcy's pride in his social position initially makes him appear arrogant and disdainful\n",
      "- Elizabeth's prejudice against Darcy is based on first impressions and incomplete information\n",
      "- Both characters must overcome these flaws to find happiness together\n",
      "\n",
      "2. Marriage and Social Class:\n",
      "- The novel explores different types of marriages (for love, security, or social advancement)\n",
      "- Social class differences and their impact on relationships\n",
      "- The economic necessity of marriage for women in the period\n",
      "- Contrast between marriages based on true understanding (Elizabeth/Darcy, Jane/Bingley) versus those based on superficial attractions (Lydia/Wickham) or practical considerations (Charlotte/Collins)\n",
      "\n",
      "3. Reputation and Social Expectations:\n",
      "- The importance of reputation, particularly for women\n",
      "- Social manners and proper behavior\n",
      "- The constraints placed on women in Regency society\n",
      "\n",
      "4. Personal Growth and Self-Knowledge:\n",
      "- Elizabeth and Darcy's journey to better understand themselves and each other\n",
      "- The importance of overcoming first impressions and preconceptions\n",
      "- Learning from mistakes and misj\n",
      "User input tokens: 4\n",
      "Output tokens: 300\n",
      "Input tokens (cache read): 187696\n",
      "Input tokens (cache write): 298\n",
      "100.0% of input prompt cached (187700 tokens)\n",
      "Time taken: 7.07 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 7-11 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": "base",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}