{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Lab 6 Answer Key*\n",
    "\n",
    "# Model Selection & Optimization\n",
    "\n",
    "**Objective:** prove to yourself that caching works, that you can break it, and that you can read the true cost off usage.\n",
    "\n",
    "Two teams can build the same feature on Claude; one pays 10x more because it defaults to the biggest model everywhere, never caches, and estimates cost with the wrong tokenizer. This lab makes the numbers concrete: measure a real cache hit, break it with the most common silent invalidator, prove the minimum-prefix rule, build a cost model from `response.usage`, and sweep effort levels on an agentic task.\n",
    "\n",
    "**Setup:** `pip install anthropic` and set `ANTHROPIC_API_KEY` in your environment. Never hardcode a key.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Measure a cache hit\n",
    "\n",
    "Caching is a byte-exact prefix match, rendered `tools` then `system` then `messages`, so a breakpoint on the last system block caches tools and system together. Verify the prompt clears Opus 4.8's 4096-token minimum with the API's own counter, then read the proof off `usage`: cold request writes, warm request reads.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import anthropic\n",
    "from datetime import datetime\n",
    "\n",
    "client = anthropic.Anthropic()  # resolves credentials from the environment\n",
    "\n",
    "# A frozen context large enough to clear Opus 4.8's 4096-token minimum\n",
    "# cacheable prefix.\n",
    "LARGE_STABLE_CONTEXT = \"\\n\".join(\n",
    "    f\"Policy {i:03d}: a severity {i % 5} incident pages the on-call rotation \"\n",
    "    f\"within {i % 30 + 1} minutes and requires a postmortem within {i % 7 + 1} days.\"\n",
    "    for i in range(500)\n",
    ")\n",
    "\n",
    "SYSTEM = [\n",
    "    {\"type\": \"text\", \"text\": \"You are an SRE assistant for Acme Cloud.\"},\n",
    "    {\n",
    "        \"type\": \"text\",\n",
    "        \"text\": LARGE_STABLE_CONTEXT,\n",
    "        \"cache_control\": {\"type\": \"ephemeral\"},   # breakpoint on the LAST block\n",
    "    },\n",
    "]\n",
    "\n",
    "# Size the prefix with the API's own counter (never tiktoken: it is\n",
    "# OpenAI's tokenizer and undercounts Claude).\n",
    "counted = client.messages.count_tokens(\n",
    "    model=\"claude-opus-4-8\",\n",
    "    system=SYSTEM,\n",
    "    messages=[{\"role\": \"user\", \"content\": \"placeholder\"}],\n",
    ")\n",
    "print(\"prompt tokens:\", counted.input_tokens, \"(must exceed 4096 to cache)\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def report(resp, label):\n",
    "    u = resp.usage\n",
    "    print(f\"{label}: write={u.cache_creation_input_tokens} \"\n",
    "          f\"read={u.cache_read_input_tokens} uncached={u.input_tokens}\")\n",
    "\n",
    "\n",
    "params = dict(\n",
    "    model=\"claude-opus-4-8\", max_tokens=512, system=SYSTEM,\n",
    "    messages=[{\"role\": \"user\",\n",
    "               \"content\": \"Who gets paged for a severity 2 incident?\"}],\n",
    ")\n",
    "r1 = client.messages.create(**params); report(r1, \"req 1\")   # cold: write > 0, read == 0\n",
    "r2 = client.messages.create(**params); report(r2, \"req 2\")   # warm: read > 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Break it deliberately\n",
    "\n",
    "The most common caching bug in production is a byte that changes every request, buried in the prefix. There is no error when caching fails, just a bill that never goes down, so reproduce the failure until you can spot it at a glance.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WRONG: datetime.now() changes the prefix bytes on EVERY request. Every\n",
    "# request writes a new cache entry; none is ever read.\n",
    "broken_system = [{\n",
    "    \"type\": \"text\",\n",
    "    \"text\": f\"You are an SRE assistant. Current time: {datetime.now()}.\\n\\n\"\n",
    "            + LARGE_STABLE_CONTEXT,\n",
    "    \"cache_control\": {\"type\": \"ephemeral\"},\n",
    "}]\n",
    "for i in (1, 2):\n",
    "    r = client.messages.create(\n",
    "        model=\"claude-opus-4-8\", max_tokens=512,\n",
    "        system=[{**broken_system[0],\n",
    "                 \"text\": f\"You are an SRE assistant. Current time: {datetime.now()}.\\n\\n\"\n",
    "                         + LARGE_STABLE_CONTEXT}],\n",
    "        messages=params[\"messages\"],\n",
    "    )\n",
    "    report(r, f\"broken req {i}\")   # read stays 0 forever\n",
    "\n",
    "# RIGHT: freeze the prefix; the timestamp rides in the user turn, AFTER\n",
    "# the breakpoint, where it invalidates nothing.\n",
    "for i in (1, 2):\n",
    "    r = client.messages.create(\n",
    "        model=\"claude-opus-4-8\", max_tokens=512,\n",
    "        system=SYSTEM,\n",
    "        messages=[{\"role\": \"user\",\n",
    "                   \"content\": f\"[time: {datetime.now()}] Who gets paged for a severity 2 incident?\"}],\n",
    "    )\n",
    "    report(r, f\"fixed req {i}\")    # req 2 reads the cache again\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Prove the minimum\n",
    "\n",
    "The minimum cacheable prefix is model-dependent: 4096 tokens on Opus 4.8. A shorter prefix does not error and does not cache; `cache_creation_input_tokens` simply stays 0 no matter how many breakpoints you add.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Roughly 3,000 tokens: below Opus 4.8's 4096-token minimum.\n",
    "SMALL_CONTEXT = \"\\n\".join(\n",
    "    f\"Policy {i:03d}: a severity {i % 5} incident pages the on-call rotation \"\n",
    "    f\"within {i % 30 + 1} minutes and requires a postmortem within {i % 7 + 1} days.\"\n",
    "    for i in range(140)\n",
    ")\n",
    "small_system = [{\n",
    "    \"type\": \"text\",\n",
    "    \"text\": \"You are an SRE assistant for Acme Cloud.\\n\\n\" + SMALL_CONTEXT,\n",
    "    \"cache_control\": {\"type\": \"ephemeral\"},   # a breakpoint that does nothing\n",
    "}]\n",
    "\n",
    "counted = client.messages.count_tokens(\n",
    "    model=\"claude-opus-4-8\", system=small_system,\n",
    "    messages=params[\"messages\"],\n",
    ")\n",
    "print(\"prompt tokens:\", counted.input_tokens, \"(below the 4096 minimum)\")\n",
    "\n",
    "for i in (1, 2):\n",
    "    r = client.messages.create(\n",
    "        model=\"claude-opus-4-8\", max_tokens=512,\n",
    "        system=small_system, messages=params[\"messages\"],\n",
    "    )\n",
    "    report(r, f\"small req {i}\")\n",
    "# write == 0 AND read == 0 on both requests, with no error: the prefix\n",
    "# is below the minimum and silently does not cache.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Build a cost model\n",
    "\n",
    "`input_tokens` is the uncached remainder only. Total prompt size is the sum of all three input fields, and the dollar cost weights each field differently: cache writes at 1.25x input price (5-minute TTL), cache reads at about 0.1x, output at the output price. Modeling cost off `input_tokens` alone makes caching look like it did nothing.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Opus 4.8: $5 in / $25 out per MTok.\n",
    "IN, OUT = 5.0 / 1e6, 25.0 / 1e6\n",
    "\n",
    "\n",
    "def request_cost(u):\n",
    "    return (\n",
    "        u.input_tokens                  * IN\n",
    "        + u.cache_creation_input_tokens * IN * 1.25   # 5-min write premium\n",
    "        + u.cache_read_input_tokens     * IN * 0.10   # cache read\n",
    "        + u.output_tokens               * OUT\n",
    "    )\n",
    "\n",
    "\n",
    "def total_prompt(u):\n",
    "    # The full prompt size is the SUM of the three input fields.\n",
    "    return (u.input_tokens + u.cache_creation_input_tokens\n",
    "            + u.cache_read_input_tokens)\n",
    "\n",
    "\n",
    "for label, resp in ((\"req 1 (cold)\", r1), (\"req 2 (warm)\", r2)):\n",
    "    print(f\"{label}: ${request_cost(resp.usage):.5f}  \"\n",
    "          f\"prompt={total_prompt(resp.usage)} tokens\")\n",
    "\n",
    "# Same prompt both times; usage just splits it differently across the\n",
    "# three input fields. The warm request is the cheaper one.\n",
    "assert total_prompt(r1.usage) == total_prompt(r2.usage)\n",
    "assert request_cost(r2.usage) < request_cost(r1.usage)\n",
    "print(\"cached rerun saved:\",\n",
    "      f\"${request_cost(r1.usage) - request_cost(r2.usage):.5f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Sweep effort levels\n",
    "\n",
    "Model tier sets the capability ceiling; `effort` (inside `output_config`, default `high`) sets how hard the model works within it. On agentic work, higher effort often consolidates tool calls, so a better plan up front can cost less in total than a low-effort run that flails through many turns.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "calc_tool = {\n",
    "    \"name\": \"calculator\",\n",
    "    \"description\": \"Evaluate a basic arithmetic expression, e.g. (2 + 3) * 4.\",\n",
    "    \"input_schema\": {\"type\": \"object\",\n",
    "                     \"properties\": {\"expression\": {\"type\": \"string\"}},\n",
    "                     \"required\": [\"expression\"]},\n",
    "}\n",
    "\n",
    "\n",
    "def safe_eval(expr):\n",
    "    # The expression is untrusted model output (Lab 7): allowlist, then eval.\n",
    "    if not re.fullmatch(r\"[0-9+\\-*/(). ]+\", expr):\n",
    "        raise ValueError(\"expression not on the arithmetic allowlist\")\n",
    "    return eval(expr, {\"__builtins__\": {}}, {})\n",
    "\n",
    "\n",
    "TASK = (\"Use the calculator tool to compute each customer's invoice, then the \"\n",
    "        \"grand total. A: 120 GB-hours at $2.00. B: 45 GB-hours at $2.00 with a \"\n",
    "        \"10 percent discount. C: 300 GB-hours at $1.50. Report the grand total.\")\n",
    "\n",
    "\n",
    "def run_at(effort):\n",
    "    msgs = [{\"role\": \"user\", \"content\": TASK}]\n",
    "    turns, out_tokens, cost = 0, 0, 0.0\n",
    "    while True:\n",
    "        r = client.messages.create(\n",
    "            model=\"claude-opus-4-8\", max_tokens=4096,\n",
    "            thinking={\"type\": \"adaptive\"},\n",
    "            output_config={\"effort\": effort},   # low | medium | high | xhigh | max\n",
    "            tools=[calc_tool], messages=msgs,\n",
    "        )\n",
    "        turns += 1\n",
    "        out_tokens += r.usage.output_tokens\n",
    "        cost += request_cost(r.usage)\n",
    "        if r.stop_reason != \"tool_use\":\n",
    "            return turns, out_tokens, cost\n",
    "        msgs.append({\"role\": \"assistant\", \"content\": r.content})\n",
    "        results = [{\"type\": \"tool_result\", \"tool_use_id\": b.id,\n",
    "                    \"content\": str(safe_eval(b.input[\"expression\"]))}\n",
    "                   for b in r.content if b.type == \"tool_use\"]\n",
    "        msgs.append({\"role\": \"user\", \"content\": results})\n",
    "\n",
    "\n",
    "for effort in (\"low\", \"high\", \"xhigh\"):\n",
    "    turns, out_tokens, cost = run_at(effort)\n",
    "    print(f\"{effort:>5}: turns={turns} output_tokens={out_tokens} cost=${cost:.5f}\")\n",
    "# Expect higher effort to consolidate tool calls into fewer turns. Fewer\n",
    "# turns means fewer resends of the growing transcript, which is why higher\n",
    "# effort can be the CHEAPER configuration on agentic work.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}