{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a056f45c",
   "metadata": {
    "id": "a056f45c"
   },
   "source": [
    "# W1 Lab — First Model Calls and Prompting Fundamentals\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ralbu85/stml_2026/blob/main/lectures/week01/W1_lab_setup.ipynb)\n",
    "\n",
    "**Goal.** Call a language model from code, know what one call does and does not do,\n",
    "and control it through the system prompt — precisely enough that the final task's\n",
    "answers come back as clean JSON, checked by machine.\n",
    "\n",
    "Why this comes first: the notes (Ch. 1) define an agent as code that acts on model\n",
    "output, so everything starts with making the call and making its output usable by code.\n",
    "\n",
    "The path: setup (1) → one call taken apart: roles, statelessness, temperature (2) →\n",
    "why \"just ask for JSON\" is not enough (3) → the core task + two exercises (4) → cost\n",
    "and completion (5) → prompt-rewriting drills (6).\n",
    "\n",
    "*Runtime:* Google Colab, top-to-bottom, ~80 minutes. Cells marked ✍️ ask for your own\n",
    "writing — a fill-in or a written prediction.\n",
    "\n",
    "*Sources:* lab format adapted from DeepLearning.AI, *Agentic AI* (Andrew Ng); prompting\n",
    "exercises draw on Anthropic's *Prompt Engineering Interactive Tutorial*.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36aa2282",
   "metadata": {
    "id": "36aa2282"
   },
   "source": [
    "## 1. Setup\n",
    "\n",
    "### 1.1 Installation\n",
    "\n",
    "`aisuite` exposes multiple providers (OpenAI, Anthropic) behind one interface, so the\n",
    "same code runs whichever provider your key belongs to.\n",
    "\n",
    "*Do:* run the cell (~30 seconds, once per session).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbdd678",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 10593,
     "status": "ok",
     "timestamp": 1788237305443,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "3fbdd678",
    "outputId": "55b71de4-958a-4d7d-f7c9-4d86980a4f87"
   },
   "outputs": [],
   "source": [
    "%pip install -q \"aisuite[openai,anthropic]\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "558f4f16",
   "metadata": {
    "id": "558f4f16"
   },
   "source": [
    "### 1.2 API key and model\n",
    "\n",
    "An **API key** = the secret string that identifies your account to the provider and\n",
    "bills usage to it (issuing steps: the API Setup guide on the course site). The key is\n",
    "yours; do not share the notebook with the key still inside.\n",
    "\n",
    "*Do:* replace `PASTE-YOUR-KEY-HERE` with your key and run the cell. Nothing prints.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4919aa8c",
   "metadata": {
    "executionInfo": {
     "elapsed": 12,
     "status": "ok",
     "timestamp": 1788237328437,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "4919aa8c"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"PASTE-YOUR-KEY-HERE\"\n",
    "\n",
    "MODEL = \"openai:gpt-4o-mini\"   # Anthropic accounts: MODEL = \"anthropic:claude-haiku-4-5\" and set ANTHROPIC_API_KEY instead"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6548a16f",
   "metadata": {
    "id": "6548a16f"
   },
   "source": [
    "### 1.3 One Call, Written Out\n",
    "\n",
    "A **message** = one turn of the exchange: a dict with a `role` and a `content` string.\n",
    "`user` carries the request. A call sends a list of messages and gets back a response\n",
    "object; the reply text sits at `response.choices[0].message.content`.\n",
    "\n",
    "*Do:* run the cell. Compare the `messages` list you sent with the reply you got back.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "977cae72",
   "metadata": {
    "id": "977cae72"
   },
   "outputs": [],
   "source": [
    "import aisuite\n",
    "\n",
    "client = aisuite.Client()\n",
    "\n",
    "messages = [{\"role\": \"user\", \"content\": \"Reply with exactly: ready\"}]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "\n",
    "print(\"messages sent:\", messages)\n",
    "print(\"reply text   :\", response.choices[0].message.content)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "782c5092",
   "metadata": {
    "id": "782c5092"
   },
   "source": [
    "Any error here is a setup problem, not a code problem — recheck the API Setup guide\n",
    "before continuing.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d672ebb3",
   "metadata": {
    "id": "d672ebb3"
   },
   "source": [
    "## 2. Anatomy of a Model Call\n",
    "\n",
    "### 2.1 Messages and roles\n",
    "\n",
    "Section 1 sent a list holding one `user` message. A message can also carry `system`,\n",
    "for standing instructions, or `assistant`, for the model's own previous replies — a\n",
    "call receives a list of these and nothing else.\n",
    "\n",
    "The two cells below ask the same question, each call written out the way Section 1\n",
    "wrote it. The two lists differ by one dict: a `system` message at the front of the\n",
    "second one.\n",
    "\n",
    "*Do:* run both cells and compare the two answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7fde95a",
   "metadata": {
    "id": "c7fde95a"
   },
   "outputs": [],
   "source": [
    "QUESTION = \"What is an AI agent?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "SP_X3KkyQYsN",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 2894,
     "status": "ok",
     "timestamp": 1788237446241,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "SP_X3KkyQYsN",
    "outputId": "810d5e56-d928-492a-825a-9c0d581be6f8"
   },
   "outputs": [],
   "source": [
    "## Without System Prompt\n",
    "plain_messages = [\n",
    "    {\"role\": \"user\", \"content\": QUESTION},\n",
    "]\n",
    "plain = client.chat.completions.create(model=MODEL, messages=plain_messages, temperature=0.0)\n",
    "print(\"\\nPLAIN:\\n\", plain.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "XNgznwToQOs3",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 712,
     "status": "ok",
     "timestamp": 1788237451058,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "XNgznwToQOs3",
    "outputId": "46f9056f-4ef0-4f85-bafa-cd413f907907"
   },
   "outputs": [],
   "source": [
    "## With System Prompt\n",
    "instructed_messages = [\n",
    "    {\"role\": \"system\", \"content\": \"Answer in one sentence, for a graduate ML audience.\"},\n",
    "    {\"role\": \"user\", \"content\": QUESTION},\n",
    "]\n",
    "instructed = client.chat.completions.create(model=MODEL, messages=instructed_messages, temperature=0.0)\n",
    "print(\"\\nINSTRUCTED:\\n\", instructed.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a424b9b6",
   "metadata": {
    "id": "a424b9b6"
   },
   "source": [
    "The instruction is a separate dict at the front of the list, not text prepended to the\n",
    "question. An application fixes that dict once and varies only the `user` entry —\n",
    "Section 4 turns this into the task.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd6e8603",
   "metadata": {
    "id": "dd6e8603"
   },
   "source": [
    "### 2.2 Statelessness\n",
    "\n",
    "An API call is **stateless** = the model sees only the message list inside that one\n",
    "call; the provider keeps no conversation state between calls. The experiment tells one\n",
    "call a fact and asks the next call to repeat it.\n",
    "\n",
    "*Do:* run both cells and read the second answer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "235c7d9e",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "executionInfo": {
     "elapsed": 968,
     "status": "ok",
     "timestamp": 1788237525736,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "235c7d9e",
    "outputId": "e8e8e885-181a-4587-a89f-16b0cbb89491"
   },
   "outputs": [],
   "source": [
    "first_messages = [\n",
    "    {\"role\": \"user\", \"content\": \"My research topic is maritime logistics optimization. \"\n",
    "                                \"Acknowledge in five words.\"},\n",
    "]\n",
    "first = client.chat.completions.create(model=MODEL, messages=first_messages, temperature=0.0)\n",
    "first_reply = first.choices[0].message.content\n",
    "first_reply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "SuRwdxkJQsSl",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "executionInfo": {
     "elapsed": 1033,
     "status": "ok",
     "timestamp": 1788237535476,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "SuRwdxkJQsSl",
    "outputId": "42f85607-05ed-4279-ca39-81e2ce0f7b98"
   },
   "outputs": [],
   "source": [
    "second_messages = [\n",
    "    {\"role\": \"user\", \"content\": \"What is my research topic?\"},\n",
    "]\n",
    "second = client.chat.completions.create(model=MODEL, messages=second_messages, temperature=0.0)\n",
    "second_reply = second.choices[0].message.content\n",
    "second_reply"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26c8203",
   "metadata": {
    "id": "a26c8203"
   },
   "source": [
    "The second call cannot answer, because the first call's content never reached it.\n",
    "Conversation is an application-side construct: the client resends the accumulated\n",
    "history with every call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc4dc623",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1900,
     "status": "ok",
     "timestamp": 1788237549175,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "cc4dc623",
    "outputId": "6520463c-59fe-4947-df0a-f7bb9160ad55"
   },
   "outputs": [],
   "source": [
    "conversation = [\n",
    "    {\"role\": \"user\", \"content\": \"My research topic is maritime logistics optimization. \"\n",
    "                                \"Acknowledge in five words.\"},\n",
    "]\n",
    "response = client.chat.completions.create(model=MODEL, messages=conversation, temperature=0.0)\n",
    "conversation.append({\"role\": \"assistant\", \"content\": response.choices[0].message.content})\n",
    "conversation.append({\"role\": \"user\", \"content\": \"What is my research topic?\"})\n",
    "\n",
    "response = client.chat.completions.create(model=MODEL, messages=conversation, temperature=0.0)\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f2a8e6e",
   "metadata": {
    "id": "1f2a8e6e"
   },
   "source": [
    "With the history resent, the model answers. Everything a model appears to remember is\n",
    "text that some code placed into the message list."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "018806c3",
   "metadata": {
    "id": "018806c3"
   },
   "source": [
    "### Exercise — one more turn ✍️\n",
    "\n",
    "Add a third user turn to `conversation` (any follow-up about the topic), resend the\n",
    "full history with another `create` call, and print the reply. The pattern is the\n",
    "`append` and `create` lines above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8471f41",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1440,
     "status": "ok",
     "timestamp": 1788237602993,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "d8471f41",
    "outputId": "1f3db0d1-2f47-4e1d-c1f5-0425e3400c14"
   },
   "outputs": [],
   "source": [
    "### FILL IN (START) ###\n",
    "# conversation.append({\"role\": \"user\", \"content\": \"...\"})\n",
    "# response = client.chat.completions.create(model=MODEL, messages=conversation, temperature=0.0)\n",
    "# print(response.choices[0].message.content)\n",
    "### FILL IN (END) ###"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc104c7e",
   "metadata": {
    "id": "bc104c7e"
   },
   "source": [
    "### 2.3 Temperature\n",
    "\n",
    "**Temperature** = the sampling parameter that scales the spread of the next-token\n",
    "distribution: 0 collapses sampling toward the most probable token, higher values admit\n",
    "less probable ones. Each cell repeats one open question five times at one setting.\n",
    "\n",
    "*Do:* run both cells and compare within each block.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0340a12",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4978,
     "status": "ok",
     "timestamp": 1788237914143,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "f0340a12",
    "outputId": "f1d55885-9b9b-45cb-8ea6-30d25177763b"
   },
   "outputs": [],
   "source": [
    "temperature = 0\n",
    "\n",
    "probe_messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Name one promising research direction for LLM agents, \"\n",
    "                                \"in one sentence.\"},\n",
    "]\n",
    "\n",
    "for i in range(5):\n",
    "    response = client.chat.completions.create(\n",
    "        model=MODEL, messages=probe_messages, temperature=temperature)\n",
    "    print(f\"t={temperature} :\", response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08l_rQ0rSNkd",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4413,
     "status": "ok",
     "timestamp": 1788237918557,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "08l_rQ0rSNkd",
    "outputId": "d35e2e77-a43c-4d18-a0b6-771d3e026973"
   },
   "outputs": [],
   "source": [
    "temperature = 1\n",
    "\n",
    "probe_messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Name one promising research direction for LLM agents, \"\n",
    "                                \"in one sentence.\"},\n",
    "]\n",
    "\n",
    "for i in range(5):\n",
    "    response = client.chat.completions.create(\n",
    "        model=MODEL, messages=probe_messages, temperature=temperature)\n",
    "    print(f\"t={temperature} :\", response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a13f02",
   "metadata": {
    "id": "d3a13f02"
   },
   "source": [
    "The five `t=0.0` runs are (near-)identical; the `t=1.0` runs differ. Deterministic\n",
    "settings suit checking and grading; sampled diversity is itself useful and W2 builds a\n",
    "technique on it (self-consistency)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63304c09",
   "metadata": {
    "id": "63304c09"
   },
   "source": [
    "## 3. Format Control\n",
    "\n",
    "Code must act on the output, and code can act only on what it can parse.\n",
    "**Structured output** = model output constrained to a machine-parseable format (here\n",
    "JSON). Models often wrap JSON in Markdown code fences, so the parser below strips them\n",
    "before `json.loads`.\n",
    "\n",
    "*Do:* run the cell, more than once if you like. A parsed dict means the loose request\n",
    "happened to work this time; `PARSE FAILED` means it did not — nothing in the request\n",
    "guarantees either.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fd582b4",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 964,
     "status": "ok",
     "timestamp": 1788237979546,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "0fd582b4",
    "outputId": "9d6910d7-af99-490c-e100-c2e0a9fdc71a"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import re\n",
    "\n",
    "def parse_json_output(text):\n",
    "    \"\"\"Model output -> parsed JSON object; raises ValueError if unparseable.\"\"\"\n",
    "    stripped = re.sub(r\"^```(?:json)?\\s*|\\s*```$\", \"\", text.strip())\n",
    "    try:\n",
    "        return json.loads(stripped)\n",
    "    except json.JSONDecodeError as exc:\n",
    "        raise ValueError(f\"unparseable model output: {text[:120]!r}\") from exc\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Give the title and year of the Chain-of-Thought paper as JSON.\"},\n",
    "]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "attempt = response.choices[0].message.content\n",
    "\n",
    "print(attempt)\n",
    "try:\n",
    "    print(parse_json_output(attempt))\n",
    "except ValueError as exc:\n",
    "    print(\"PARSE FAILED:\", exc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9556c6",
   "metadata": {
    "id": "7d9556c6"
   },
   "source": [
    "A loose request sometimes yields prose around the object. The fix is not post-hoc\n",
    "string surgery; it is an instruction that pins the format exactly — Section 4.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244dadf7",
   "metadata": {
    "id": "244dadf7"
   },
   "source": [
    "## 4. Task — Answers as JSON ✍️ (core)\n",
    "\n",
    "One system prompt, stated precisely enough that its effect passes a machine check.\n",
    "\n",
    "An example first: the cell below solves the same problem for a smaller schema — a\n",
    "sentiment classifier that must return `{\"label\": ..., \"confidence\": ...}`. Read the\n",
    "prompt before running it: it names every key, states the closed vocabulary for\n",
    "`label`, and forbids everything else.\n",
    "\n",
    "In the message list, the schema lives entirely in the `system`\n",
    "turn, and the `user` turn carries only the text to classify.\n",
    "\n",
    "*Do:* run the cell and read the raw reply next to its parsed form.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "581b8127",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 494,
     "status": "ok",
     "timestamp": 1788238397531,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -540
    },
    "id": "581b8127",
    "outputId": "1b7bf80d-210f-4fec-81fa-ff2666e3bc35"
   },
   "outputs": [],
   "source": [
    "EXAMPLE_SYSTEM = (\n",
    "    '''\n",
    "    You are a sentiment classifier. For every user text, return ONLY a raw JSON object —\n",
    "    no prose, no Markdown code fences — with exactly these keys and no others:\n",
    "\n",
    "    \"label\": exactly one of \"positive\", \"neutral\", \"negative\";\n",
    "    \"confidence\": a number between 0 and 1.\n",
    "\n",
    "    Never add keys, comments, or explanations outside the JSON object.\n",
    "    '''\n",
    ")\n",
    "EXAMPLE_USER = \"Classify the sentiment of this text.\\n\\nText: {text}\"\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": EXAMPLE_SYSTEM},\n",
    "    {\"role\": \"user\", \"content\": EXAMPLE_USER.format(text=\"The lab finally makes sense to me.\")},\n",
    "]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "reply = response.choices[0].message.content\n",
    "\n",
    "print(\"raw reply:\", reply)\n",
    "print(\"parsed:   \", parse_json_output(reply))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4b249a3",
   "metadata": {
    "id": "e4b249a3"
   },
   "source": [
    "### 4.1 Writing `JSON_SYSTEM` ✍️\n",
    "\n",
    "Write `JSON_SYSTEM` so that the answer to any question comes back as raw JSON with\n",
    "exactly these keys:\n",
    "\n",
    "| key | content |\n",
    "|---|---|\n",
    "| `topic` | short label for the subject area |\n",
    "| `answer` | one-sentence direct answer |\n",
    "| `difficulty` | one of `\"intro\"`, `\"intermediate\"`, `\"advanced\"` |\n",
    "\n",
    "Follow the pattern of `EXAMPLE_SYSTEM`: name every key, state the closed vocabulary\n",
    "for `difficulty`, forbid everything else — no prose around the object, no code fences,\n",
    "no extra keys.\n",
    "\n",
    "*Do:* fill in `JSON_SYSTEM`, run the check cell, iterate until both questions read\n",
    "`PASS`. Typical first failures: code fences, or extra keys the model volunteers —\n",
    "both disappear once forbidden explicitly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac063f9",
   "metadata": {
    "id": "5ac063f9"
   },
   "outputs": [],
   "source": [
    "### FILL IN (START) ###\n",
    "JSON_SYSTEM = (\n",
    "    \"\"\n",
    ")\n",
    "### FILL IN (END) ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de171971",
   "metadata": {
    "id": "de171971"
   },
   "outputs": [],
   "source": [
    "QUESTIONS = [\n",
    "    \"How can I use an LLM to answer questions over our lab's PDF reports?\",\n",
    "    \"Why does asking the model to think step by step improve accuracy?\",\n",
    "]\n",
    "REQUIRED_KEYS = {\"topic\", \"answer\", \"difficulty\"}\n",
    "\n",
    "n_passed = 0\n",
    "for question in QUESTIONS:\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": JSON_SYSTEM},\n",
    "        {\"role\": \"user\", \"content\": question},\n",
    "    ]\n",
    "    response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "    output = response.choices[0].message.content\n",
    "    try:\n",
    "        ok = set(parse_json_output(output)) == REQUIRED_KEYS\n",
    "    except ValueError:\n",
    "        ok = False\n",
    "    n_passed += ok\n",
    "    print(f\"{'PASS' if ok else 'FAIL'}  {question[:55]}...\")\n",
    "    if not ok:\n",
    "        print(\"      output began:\", repr(output[:90]))\n",
    "\n",
    "print(\"ALL PASS\" if n_passed == len(QUESTIONS) else \"KEEP ITERATING\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e07337fb",
   "metadata": {
    "id": "e07337fb"
   },
   "source": [
    "### 4.2 Exercise — temperature vs. format ✍️\n",
    "\n",
    "Prediction, written down before running: at `t=1.0`, does the answer still parse as\n",
    "valid JSON with the three keys on every run?\n",
    "\n",
    "*Do:* run the cell and compare with your prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c24eeef",
   "metadata": {
    "id": "0c24eeef"
   },
   "outputs": [],
   "source": [
    "for run in range(3):\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": JSON_SYSTEM},\n",
    "        {\"role\": \"user\", \"content\": QUESTIONS[0]},\n",
    "    ]\n",
    "    response = client.chat.completions.create(model=MODEL, messages=messages, temperature=1.0)\n",
    "    try:\n",
    "        ok = set(parse_json_output(response.choices[0].message.content)) == REQUIRED_KEYS\n",
    "    except ValueError:\n",
    "        ok = False\n",
    "    print(f\"run {run}: {'PASS' if ok else 'FAIL'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b59aa408",
   "metadata": {
    "id": "b59aa408"
   },
   "source": [
    "### 4.3 Exercise — provider-enforced JSON mode\n",
    "\n",
    "OpenAI can enforce JSON at the API level: `response_format={\"type\": \"json_object\"}`\n",
    "constrains decoding so the reply always parses. It guarantees parseable JSON, not your\n",
    "three keys — those still come from `JSON_SYSTEM` — and it is OpenAI-specific. (The API\n",
    "requires the word \"JSON\" somewhere in the messages; `JSON_SYSTEM` already satisfies\n",
    "this.)\n",
    "\n",
    "Prediction: with the parameter added, can the reply still violate the three-key schema\n",
    "even though it always parses?\n",
    "\n",
    "*Do:* run the cell.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "783d9b5f",
   "metadata": {
    "id": "783d9b5f"
   },
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": JSON_SYSTEM},\n",
    "    {\"role\": \"user\", \"content\": QUESTIONS[0]},\n",
    "]\n",
    "response = client.chat.completions.create(\n",
    "    model=MODEL, messages=messages, temperature=0.0,\n",
    "    response_format={\"type\": \"json_object\"})\n",
    "output = response.choices[0].message.content\n",
    "\n",
    "print(output)\n",
    "print(parse_json_output(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "237bed0b",
   "metadata": {
    "id": "237bed0b"
   },
   "source": [
    "## 5. Cost and Completion\n",
    "\n",
    "Every response object carries a `usage` field: the token counts the provider just\n",
    "billed for that call. The cell reads it off the JSON-mode response from 4.3 and\n",
    "prices that one call at gpt-4o-mini list prices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1155fa6d",
   "metadata": {
    "id": "1155fa6d"
   },
   "outputs": [],
   "source": [
    "PRICE_PER_M_PROMPT = 0.15        # USD per 1M input tokens, gpt-4o-mini\n",
    "PRICE_PER_M_COMPLETION = 0.60    # USD per 1M output tokens, gpt-4o-mini\n",
    "\n",
    "usage = response.usage\n",
    "cost = (usage.prompt_tokens * PRICE_PER_M_PROMPT\n",
    "        + usage.completion_tokens * PRICE_PER_M_COMPLETION) / 1_000_000\n",
    "print(f\"prompt tokens: {usage.prompt_tokens}   completion tokens: {usage.completion_tokens}\")\n",
    "print(f\"cost of that one call: ${cost:.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26c32699",
   "metadata": {
    "id": "26c32699"
   },
   "source": [
    "Multiplied over this session's few dozen calls, the lab costs on the order of a cent.\n",
    "Call and token counts, not dollars, become the binding constraint once loops multiply\n",
    "calls (notes Ch. 12).\n",
    "\n",
    "### Completion check\n",
    "\n",
    "All rows must read `PASS` before submission; grading checks these structural facts,\n",
    "never prose quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28562b4c",
   "metadata": {
    "id": "28562b4c"
   },
   "outputs": [],
   "source": [
    "completion = {\n",
    "    \"JSON-mode reply parsed\":             isinstance(output, str) and len(output) > 0,\n",
    "    \"JSON_SYSTEM written (>= 40 chars)\":  len(JSON_SYSTEM.strip()) >= 40,\n",
    "    \"both JSON questions PASS\":           n_passed == len(QUESTIONS),\n",
    "}\n",
    "for item, ok in completion.items():\n",
    "    print(f\"{'PASS' if ok else 'FAIL':4}  {item}\")\n",
    "print(\"\\nLAB COMPLETE\" if all(completion.values()) else \"\\nNOT COMPLETE YET\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "656efaa0",
   "metadata": {},
   "source": [
    "## 6. Practice — Rewriting Prompts ✍️\n",
    "\n",
    "The system prompts so far were given. These three drills hand the prompt over: in\n",
    "each, the surrounding code stays fixed and only the prompt string changes, until\n",
    "every check prints `PASS`. The first wording rarely passes — iterating on the words\n",
    "is the exercise. The drills are practice; they are not part of the Section 5\n",
    "completion check."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "929b6a0c",
   "metadata": {},
   "source": [
    "### 6.1 Word budget ✍️\n",
    "\n",
    "`TEN_WORD_SYSTEM` must make the answer to any question come back as exactly ten\n",
    "words. The starter is too soft to do it.\n",
    "\n",
    "*Do:* edit the string, run, iterate until both lines read `PASS`.\n",
    "\n",
    "Hints: state the rule as a hard constraint, not a preference; tell the model to\n",
    "count the words; say what to drop (detail, hedges) when the answer wants to run\n",
    "long."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "717c155e",
   "metadata": {},
   "outputs": [],
   "source": [
    "### FILL IN (START) ###\n",
    "TEN_WORD_SYSTEM = \"Answer briefly.\"\n",
    "### FILL IN (END) ###\n",
    "\n",
    "for q in [\"What is an AI agent?\",\n",
    "          \"Why do LLM answers vary between runs?\"]:\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": TEN_WORD_SYSTEM},\n",
    "        {\"role\": \"user\", \"content\": q},\n",
    "    ]\n",
    "    response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "    reply = response.choices[0].message.content\n",
    "    n_words = len(reply.split())\n",
    "    print(f\"{'PASS' if n_words == 10 else 'FAIL'} ({n_words:2d} words)  {reply}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbcb85b4",
   "metadata": {},
   "source": [
    "### 6.2 Audience shift ✍️\n",
    "\n",
    "Section 2.1's system message aimed the answer at a graduate ML audience.\n",
    "`AUDIENCE_SYSTEM` must aim the same `QUESTION` at a twelve-year-old: at most two\n",
    "sentences, and none of the words in `JARGON`.\n",
    "\n",
    "*Do:* edit the string, run, iterate until both checks read `PASS`.\n",
    "\n",
    "Hints: name the audience; cap the sentence count inside the instruction; ban the\n",
    "jargon list explicitly — the model cannot avoid words it was never told to avoid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f8854a",
   "metadata": {},
   "outputs": [],
   "source": [
    "### FILL IN (START) ###\n",
    "AUDIENCE_SYSTEM = \"Answer the question.\"\n",
    "### FILL IN (END) ###\n",
    "\n",
    "JARGON = [\"autonomous\", \"orchestration\", \"paradigm\", \"LLM\"]\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": AUDIENCE_SYSTEM},\n",
    "    {\"role\": \"user\", \"content\": QUESTION},\n",
    "]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "reply = response.choices[0].message.content\n",
    "\n",
    "n_sentences = len([s for s in re.split(r\"[.!?]+\", reply) if s.strip()])\n",
    "used = [w for w in JARGON if w.lower() in reply.lower()]\n",
    "print(reply)\n",
    "print(f\"{'PASS' if n_sentences <= 2 else 'FAIL'}  sentences: {n_sentences}\")\n",
    "print(f\"{'PASS' if not used else 'FAIL'}  jargon used: {used or 'none'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6600e7cd",
   "metadata": {},
   "source": [
    "### 6.3 Refusal to guess ✍️\n",
    "\n",
    "The two check cells ask about `NOTE`. The first answer is in the text; the second is\n",
    "not there at all, and under the starter prompt the model invents one anyway.\n",
    "`EXTRACT_SYSTEM` must make the model answer from the note alone and reply with\n",
    "exactly `UNKNOWN` when the note does not contain the answer.\n",
    "\n",
    "*Do:* edit the string, then run both check cells; iterate until both read `PASS`.\n",
    "\n",
    "Hints: state both rules — answer only from the text, and the exact token to return\n",
    "when the answer is absent. Models default to being helpful, so the `UNKNOWN` case\n",
    "must be demanded as required behavior, not offered as an option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e62df36",
   "metadata": {},
   "outputs": [],
   "source": [
    "NOTE = (\"Lab meeting moved to Thursday 4pm in room B203. \"\n",
    "        \"Bring the Q3 shipping-delay figures.\")\n",
    "\n",
    "### FILL IN (START) ###\n",
    "EXTRACT_SYSTEM = \"Answer questions about the note.\"\n",
    "### FILL IN (END) ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6bb1eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Answer present in the note\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": EXTRACT_SYSTEM},\n",
    "    {\"role\": \"user\", \"content\": f\"Note: {NOTE}\\n\\nQuestion: Where is the meeting?\"},\n",
    "]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "reply = response.choices[0].message.content\n",
    "print(f\"{'PASS' if 'B203' in reply else 'FAIL'}  {reply}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a89e01a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Answer absent from the note — an unguarded model invents one\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": EXTRACT_SYSTEM},\n",
    "    {\"role\": \"user\", \"content\": f\"Note: {NOTE}\\n\\nQuestion: Who chairs the meeting?\"},\n",
    "]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages, temperature=0.0)\n",
    "reply = response.choices[0].message.content\n",
    "print(f\"{'PASS' if reply.strip() == 'UNKNOWN' else 'FAIL'}  {reply}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a352b0d1",
   "metadata": {
    "id": "a352b0d1"
   },
   "source": [
    "---\n",
    "\n",
    "W2 works on the content of the prompt rather than its format: chain-of-thought,\n",
    "worked-example exemplars, and sampling-based self-consistency, measured on a small\n",
    "code-graded evalset. After a week of writing every call out by hand, W2 also\n",
    "introduces the `ask` helper that wraps the repeated lines. Reference answers for\n",
    "this lab: `labs/checkpoints/week01/solution.py`, published after the session."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [
    {
     "file_id": "https://github.com/ralbu85/stml_2026/blob/main/lectures/week01/W1_lab_setup.ipynb",
     "timestamp": 1788238576797
    }
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}