{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c00",
   "metadata": {},
   "source": [
    "# W5 Lab — Reflection and Evaluation with `Agent` Objects\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/week05/W5_lab_reflection_evals.ipynb)\n",
    "\n",
    "**Goal.** Let one agent write a short science essay, evaluate it, and improve it (reflection). Then measure answers with code and with an LLM judge (evaluation).\n",
    "\n",
    "Every model in this lab is an `Agent` object, the class from the practice notebook (v2). In Part 1, one object does all the steps of reflection: its system message gives the role and the steps, and its tool is the Wikipedia search from the practice notebook. Week 6 divides the same work among several objects.\n",
    "\n",
    "| Part | Question |\n",
    "|---|---|\n",
    "| 1. Reflection in one agent | How does one agent write, evaluate, and improve an essay, and when does it stop? |\n",
    "| 2. Evaluation with code | How often is an invoice's due date extracted correctly, and did a change help? |\n",
    "| 3. An LLM judge | How do we grade essays with a rubric, can we trust the judge, and did the reflection help? |\n",
    "\n",
    "**Cell labels.** `# ✍️ TYPE` = write the code between the `FILL IN` lines yourself; `# ✍️ WRITE` = answer in words; `# ▶ RUN · example` = read, then run; `# ▶ RUN · setup` / `check` = just run.\n",
    "\n",
    "About 80 minutes. Not collected. The search results come from Wikipedia and the model replies come from your own run, so your text will differ from the examples. The invoices in Part 2 and the graded essays in Part 3 are fictional examples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c01",
   "metadata": {},
   "source": [
    "## 0. Setup\n",
    "\n",
    "*Do:* run the four cells; paste your API key into the second. The last cell defines `show_messages` and the `Agent` class from the practice notebook (v2), with two additions for this week, both for measuring: an agent made with `Agent(system, json_mode=True)` replies in valid JSON, and `call_model` uses `temperature=0`, so the same input gives (nearly) the same reply each run. `show_messages` also prints text that arrives together with tool requests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · setup: just run\n",
    "%pip install -q \"openai==1.109.1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c03",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · setup: just run\n",
    "import os\n",
    "\n",
    "OPENAI_API_KEY = \"PASTE-YOUR-KEY-HERE\"   # paste your key between the quotes\n",
    "if OPENAI_API_KEY != \"PASTE-YOUR-KEY-HERE\":\n",
    "    os.environ[\"OPENAI_API_KEY\"] = OPENAI_API_KEY   # the client reads the key from here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · setup: just run\n",
    "import json\n",
    "from datetime import datetime\n",
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI()\n",
    "MODEL = \"gpt-4.1\"\n",
    "done = {}                          # the check cells record your results here\n",
    "messages = [{\"role\": \"user\", \"content\": \"Reply with exactly: ready\"}]\n",
    "response = client.chat.completions.create(model=MODEL, messages=messages)\n",
    "reply = response.choices[0].message.model_dump(exclude_none=True)\n",
    "messages.append(reply)\n",
    "reply[\"content\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · setup: just run\n",
    "# The course helpers and class, as in the practice notebook (v2), with two additions for this week: json_mode and temperature=0.\n",
    "import inspect\n",
    "\n",
    "PYTHON_TO_JSON = {str: \"string\", int: \"integer\", float: \"number\", bool: \"boolean\"}\n",
    "\n",
    "def tool_from_function(function):\n",
    "    \"\"\"Build a tool schema from a function's name, docstring, and type annotations.\"\"\"\n",
    "    doc = inspect.getdoc(function)\n",
    "    arg_help = {}\n",
    "    for line in doc.split(\"\\n\"):\n",
    "        line = line.strip()\n",
    "        if \":\" in line and not line.endswith(\":\"):\n",
    "            name, text = line.split(\":\", 1)\n",
    "            arg_help[name.strip()] = text.strip()\n",
    "    properties = {}\n",
    "    for name, parameter in inspect.signature(function).parameters.items():\n",
    "        properties[name] = {\"type\": PYTHON_TO_JSON[parameter.annotation], \"description\": arg_help[name]}\n",
    "    return {\"type\": \"function\",\n",
    "            \"function\": {\"name\": function.__name__, \"description\": doc.split(\"\\n\")[0],\n",
    "                         \"parameters\": {\"type\": \"object\", \"properties\": properties, \"required\": list(properties)}}}\n",
    "\n",
    "\n",
    "def show_messages(messages):\n",
    "    \"\"\"Print each message: position, role, and content.\"\"\"\n",
    "    for i in range(len(messages)):\n",
    "        m = messages[i]\n",
    "        if \"tool_calls\" in m:                          # a request: print each requested function\n",
    "            if \"content\" in m and m[\"content\"] != \"\":  # text that arrived together with the requests\n",
    "                print(i, m[\"role\"], m[\"content\"])\n",
    "            for request in m[\"tool_calls\"]:\n",
    "                print(i, m[\"role\"], request[\"function\"])\n",
    "        else:\n",
    "            print(i, m[\"role\"], m[\"content\"])\n",
    "    print()\n",
    "\n",
    "\n",
    "class Agent:\n",
    "    \"\"\"A model with its own message list and, optionally, tools.\"\"\"\n",
    "\n",
    "    def __init__(self, system, tool_functions=(), json_mode=False):\n",
    "        self.messages = [{\"role\": \"system\", \"content\": system}]              # SYSTEM message, once\n",
    "        self.tools = []                                                      # TOOL SCHEMAS\n",
    "        self.functions = {}                                                  # name -> Python function\n",
    "        for function in tool_functions:                                      # () = no tools\n",
    "            self.tools.append(tool_from_function(function))\n",
    "            self.functions[function.__name__] = function\n",
    "        self.response_format = {\"type\": \"text\"}\n",
    "        if json_mode:                                                        # NEW this week: every reply is valid JSON\n",
    "            self.response_format = {\"type\": \"json_object\"}\n",
    "        # NEW this week: temperature=0 in call_model, so the same input gives (nearly) the same reply each run\n",
    "\n",
    "    def call_model(self):\n",
    "        \"\"\"Call the model with the list; keep the reply and return it.\"\"\"\n",
    "        if len(self.tools) > 0:\n",
    "            response = client.chat.completions.create(model=MODEL, messages=self.messages, tools=self.tools,\n",
    "                                                      response_format=self.response_format, temperature=0)\n",
    "        else:                                                                # no tools: the API rejects an empty list\n",
    "            response = client.chat.completions.create(model=MODEL, messages=self.messages,\n",
    "                                                      response_format=self.response_format, temperature=0)\n",
    "        reply = response.choices[0].message.model_dump(exclude_none=True)\n",
    "        self.messages.append(reply)                                          # ASSISTANT message\n",
    "        return reply\n",
    "\n",
    "    def run_tools(self, reply):\n",
    "        \"\"\"Run each request in the reply and add each result as a tool message.\"\"\"\n",
    "        for request in reply[\"tool_calls\"]:\n",
    "            name = request[\"function\"][\"name\"]\n",
    "            try:\n",
    "                args = json.loads(request[\"function\"][\"arguments\"])\n",
    "                result = self.functions[name](**args)\n",
    "            except Exception as error:                                       # unknown name or wrong arguments\n",
    "                result = \"error: \" + str(error)                              # the error becomes the result\n",
    "            self.messages.append({\"role\": \"tool\", \"tool_call_id\": request[\"id\"], \"content\": json.dumps(result)})   # TOOL message\n",
    "\n",
    "    def ask(self, question, max_steps=6):\n",
    "        \"\"\"Add the question, then loop until the model answers in text.\"\"\"\n",
    "        self.messages.append({\"role\": \"user\", \"content\": question})          # USER message\n",
    "        for step in range(max_steps):\n",
    "            reply = self.call_model()\n",
    "            if \"tool_calls\" in reply:                                        # requests: run the tools\n",
    "                self.run_tools(reply)\n",
    "            else:                                                            # a text answer: done\n",
    "                return reply[\"content\"]\n",
    "        print(\"Stopped: no text answer after\", max_steps, \"steps\")\n",
    "        return None                                                          # the step limit ended the loop\n",
    "\n",
    "    def show(self):\n",
    "        \"\"\"Print this object's messages.\"\"\"\n",
    "        show_messages(self.messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c06",
   "metadata": {},
   "source": [
    "## 1. Reflection in one agent\n",
    "\n",
    "The task is a short science essay that cites Wikipedia. One `Agent` object does the whole task. Its system message gives the role, and its tool is `wikipedia_search`.\n",
    "\n",
    "### 1.1 The search tool\n",
    "\n",
    "`wikipedia_search` is the function from the practice notebook. It returns the titles and short snippets of the top three articles.\n",
    "\n",
    "*Do:* run the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c07",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "import requests\n",
    "import html\n",
    "\n",
    "def wikipedia_search(query: str):\n",
    "    \"\"\"Search Wikipedia and return the top three results: title and a short snippet.\n",
    "\n",
    "    Args:\n",
    "        query: Search words, for example \"black hole\".\n",
    "    \"\"\"\n",
    "    try:\n",
    "        data = requests.get(\"https://en.wikipedia.org/w/api.php\",\n",
    "                            params={\"action\": \"query\", \"list\": \"search\", \"srsearch\": query,\n",
    "                                    \"srlimit\": 3, \"format\": \"json\"},\n",
    "                            headers={\"User-Agent\": \"STML-course-lab/1.0\"}, timeout=10).json()\n",
    "        lines = []\n",
    "        for item in data[\"query\"][\"search\"]:\n",
    "            snippet = item[\"snippet\"].replace('<span class=\"searchmatch\">', \"\").replace(\"</span>\", \"\")   # drop highlight tags\n",
    "            lines.append(item[\"title\"] + \" | \" + html.unescape(snippet))                               # &#039; -> '\n",
    "        return \"\\n\".join(lines)\n",
    "    except Exception as error:                     # a failed search is returned as text, not raised\n",
    "        return \"error: \" + str(error)\n",
    "\n",
    "print(wikipedia_search(\"black hole\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08",
   "metadata": {},
   "source": [
    "### 1.2 A writer without reflection\n",
    "\n",
    "The system message gives the role: a science writer who searches Wikipedia and cites the articles as `[article title]`. The question arrives in the user message.\n",
    "\n",
    "*Do:* run the cell and read the list: the searches, the results, and the essay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "QUESTION = \"How do black holes form?\"\n",
    "TASK = f\"Write a short essay on this question: {QUESTION}\"\n",
    "WRITER_SYSTEM = \"\"\"You are a science writer. Use wikipedia_search to find facts. Cite each Wikipedia article that you use as [article title].\"\"\"\n",
    "\n",
    "writer = Agent(WRITER_SYSTEM, [wikipedia_search])\n",
    "direct_essay = writer.ask(TASK)\n",
    "writer.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c10",
   "metadata": {},
   "source": [
    "**Reading the output.** The agent searched Wikipedia first and then wrote the essay, with the article titles in brackets. The loop of `ask` ended at the first text answer, as in Chapter 4: nothing read the essay before it was returned.\n",
    "\n",
    "### 1.3 The same writer, with reflection in its system message\n",
    "\n",
    "`REFLECT_SYSTEM` is `WRITER_SYSTEM` plus the steps of reflection: write a draft, evaluate it, improve it, and repeat. The agent takes one step in each reply, so each step is one message in its list. After each reply, the code sends only `Continue.`: the system message, not the code, says what the next step is.\n",
    "\n",
    "*Do:* run the three cells in order. After each one, read the new messages at the end of the list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "## Unit 1 - the first reply\n",
    "REFLECT_SYSTEM = f\"\"\"{WRITER_SYSTEM}\n",
    "Work in steps, one step in each reply: write a draft, evaluate the draft, improve the essay, and evaluate again.\n",
    "Repeat until an evaluation finds nothing to improve. Then reply with the final essay under the heading FINAL ESSAY.\"\"\"\n",
    "\n",
    "agent = Agent(REFLECT_SYSTEM, [wikipedia_search])\n",
    "reply = agent.ask(TASK)\n",
    "draft = reply                                      # the first reply, kept for Part 3\n",
    "agent.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "## Unit 2 - the next step\n",
    "reply = agent.ask(\"Continue.\")\n",
    "agent.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "## Unit 3 - the step after that\n",
    "reply = agent.ask(\"Continue.\")\n",
    "agent.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c14",
   "metadata": {},
   "source": [
    "**Reading the output.** Each reply is one step. In our run, the first reply was a draft that the model wrote from its own knowledge, without a search. The second reply was an evaluation: it listed strengths and points to improve, and one point was that the citations were missing. In the third step, the agent searched Wikipedia and then wrote an improved essay with citations. The search results come from outside the model, so they are **external feedback**: they add facts to the list before the next version. Your run can divide the steps differently; read what each reply holds.\n",
    "\n",
    "### 1.4 ✍️ Run the steps to the end\n",
    "\n",
    "The agent decides when the essay is finished: its system message tells it to write `FINAL ESSAY` then. The code only repeats `Continue.` until that heading arrives. The limit of 8 rounds stops the loop if the heading never comes.\n",
    "\n",
    "The code looks for the heading on its own line, `\"FINAL ESSAY\\n\"` (`\\n` is the end of a line), because the agent can also use the words in a sentence, for example \"Next, I will present the FINAL ESSAY.\"\n",
    "\n",
    "*Do:* write the core lines of the loop: stop when `reply` contains `\"FINAL ESSAY\\n\"`; if not, send `Continue.` and keep the new reply. The check cell prints the first line of every reply, as a map of the steps, and then the final essay."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ✍️ TYPE · exercise: write the code between the FILL IN lines yourself\n",
    "for step in range(8):                              # a limit, in case FINAL ESSAY never comes\n",
    "    ### FILL IN (START) ###\n",
    "    break                                          # replace: stop if reply contains \"FINAL ESSAY\\n\"; if not, ask the agent to continue\n",
    "    ### FILL IN (END) ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · check: just run\n",
    "for i in range(len(agent.messages)):\n",
    "    m = agent.messages[i]\n",
    "    if m[\"role\"] == \"assistant\" and \"tool_calls\" not in m:\n",
    "        print(i, m[\"content\"].split(\"\\n\")[0])      # the first line of each reply: a map of the steps\n",
    "final_essay = reply.split(\"FINAL ESSAY\\n\")[-1]     # the text after the heading\n",
    "print()\n",
    "print(final_essay)\n",
    "done[\"1.4\"] = \"FINAL ESSAY\\n\" in reply\n",
    "done[\"1.4\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c17",
   "metadata": {},
   "source": [
    "**Reading the output.** The map shows reflection in one list: a draft, an evaluation, an improved essay, another evaluation, and at the end the final essay. After each evaluation, the agent goes back to the role in its system message, the writer: the next reply is an essay again, with its citations. The agent itself ended the work when it wrote the heading `FINAL ESSAY`. If your map has no `FINAL ESSAY` line, the limit of 8 rounds ended the loop, and `final_essay` holds the last reply.\n",
    "\n",
    "### 1.5 The whole record\n",
    "\n",
    "*Do:* run the cell. Find the first evaluation and the essay after it. Which points of the evaluation did the agent take?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "agent.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c19",
   "metadata": {},
   "source": [
    "**Reading the output.** Compare the first evaluation with the essay after it. In our run, the evaluation asked for citations and for more detail, and the next version has both. Then read the last evaluation: the agent stops when an evaluation finds nothing more to improve. All steps are in one list, so every call sends the whole history again: each draft, each evaluation, and each search result. Note also who evaluates: the same object, with the same system message and the same list, wrote the essay and evaluated it. Week 6 gives the evaluation to a separate agent. Part 3 measures whether the final essay is better than the essay of 1.2, and what the steps cost."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c20",
   "metadata": {},
   "source": [
    "## 2. Evaluation with code\n",
    "\n",
    "Part 1 improved a single answer. **Evaluation** measures a system against explicit criteria over a set of cases. Our **evaluation set** is the ten invoices with their **reference answers**; our criterion is an exact match in `YYYY/MM/DD`. **Accuracy** = correct cases ÷ cases tested.\n",
    "\n",
    "The system under test is an **extractor** agent: it reads an invoice and returns its **payment due date**. Its system message fixes the output as JSON with one field, `due_date`, and the agent is made with `json_mode=True`, so every reply is valid JSON. Code reads the field and compares it with the reference. A free-text reply such as \"The invoice is due on 15 October 2026.\" would fail the exact match because of its form, not its date; the structured output removes that problem, so the grade depends on the date only.\n",
    "\n",
    "Every case gets a **new** `Agent` object, so its list starts empty. One object reused for all cases would carry earlier invoices and answers into later ones."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c21",
   "metadata": {},
   "source": [
    "### The invoices and their reference dates\n",
    "\n",
    "`INVOICES` holds ten invoice texts; `REFERENCES` holds the correct due date of each, worked out by a person. The references are **never sent to the model**; we use them only to check answers.\n",
    "\n",
    "*Do:* run the cell and read invoices H, I, and J."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "INVOICES = {\n",
    "    \"A\": \"Invoice A. Issued: 2026/10/01. Due: 2026/10/15. Amount: 420 USD.\",\n",
    "    \"B\": \"Invoice B. Date of issue: 3 November 2026. Payment terms: net 30 days from the date of issue. Amount: 1,250 USD.\",\n",
    "    \"C\": \"Invoice C. Issued 2026-11-10. Goods delivered 2026-11-20. Payment is due 14 days after delivery.\",\n",
    "    \"D\": \"Invoice D. Issue date: 2026/11/05. Original due date: 2026/12/01. By an agreement dated 2026/11/25, the due date is extended to 2026/12/15.\",\n",
    "    \"E\": \"Invoice E. Issued on 1 December 2026. Please pay by December 20th, 2026. A 2% late fee applies to payments received after 27 December 2026.\",\n",
    "    \"F\": \"Invoice F. Issued: 15/12/2026. Payment due: 05/01/2027. All dates are written day/month/year.\",\n",
    "    \"G\": \"Invoice G. Issued 2026/10/15. Terms: net 30 days. If the due date falls on a Saturday or Sunday, payment is due on the following Monday.\",\n",
    "    \"H\": \"Invoice H. Issued 2027/02/05. Payment is due on the last day of the month in which the invoice was issued.\",\n",
    "    \"I\": \"Invoice I. Goods delivered on 2026/12/22. Payment is due 10 business days after delivery. Business days are Monday to Friday, except the holidays 2026/12/25 and 2027/01/01.\",\n",
    "    \"J\": \"Invoice J. Issued 2026/09/14. Payment is due on the first Friday that falls at least 60 days after the issue date.\",\n",
    "}\n",
    "REFERENCES = {\n",
    "    \"A\": \"2026/10/15\", \"B\": \"2026/12/03\", \"C\": \"2026/12/04\", \"D\": \"2026/12/15\", \"E\": \"2026/12/20\",\n",
    "    \"F\": \"2027/01/05\", \"G\": \"2026/11/16\", \"H\": \"2027/02/28\", \"I\": \"2027/01/07\", \"J\": \"2026/11/13\",\n",
    "}\n",
    "print(INVOICES[\"J\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c23",
   "metadata": {},
   "source": [
    "### 2.1 The check\n",
    "\n",
    "*Do:* run the cell. The first line is `False`: a valid-looking date can still be the wrong date."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "def date_matches(pred, ref):\n",
    "    \"\"\"True when the predicted date string equals the reference exactly.\"\"\"\n",
    "    return pred == ref\n",
    "\n",
    "print(date_matches(\"2026/10/01\", \"2026/10/15\"))   # the issue date, not the due date\n",
    "print(date_matches(\"2026/10/15\", \"2026/10/15\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c25",
   "metadata": {},
   "source": [
    "### 2.2 The extractor's accuracy\n",
    "\n",
    "*Do:* run the cell. It prints one line per invoice (result, reference, prediction), the accuracy, and the last extractor's list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "EXTRACT_SYSTEM = \"\"\"You read an invoice and find its payment due date.\n",
    "Return JSON: {\"due_date\": \"YYYY/MM/DD\"}\"\"\"\n",
    "\n",
    "correct = 0\n",
    "for name in INVOICES:                                     # the keys: \"A\", \"B\", ...\n",
    "    extractor = Agent(EXTRACT_SYSTEM, json_mode=True)      # a new object for every case\n",
    "    answer = json.loads(extractor.ask(INVOICES[name]))     # JSON text -> dictionary\n",
    "    pred = answer[\"due_date\"]\n",
    "    ok = date_matches(pred, REFERENCES[name])\n",
    "    if ok:\n",
    "        correct = correct + 1\n",
    "    print(name, ok, REFERENCES[name], pred)\n",
    "print(\"accuracy:\", correct, \"/\", len(INVOICES))\n",
    "extractor.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c27",
   "metadata": {},
   "source": [
    "### 2.3 A change on the same cases: the extractor checks its answer\n",
    "\n",
    "Same invoices, same check; one part changes. After its first answer, the same extractor gets one more message: check the answer and return the JSON again. This is reflection in one agent, as in Part 1, but without a tool: no new information arrives between the answer and the check. Compare the accuracy **and** the number of calls with 2.2.\n",
    "\n",
    "*Do:* run the cell. It prints the first answer (`draft`) and the answer after the check (`final`) for each invoice, and then the list of the last extractor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "correct = 0\n",
    "calls = 0\n",
    "for name in INVOICES:\n",
    "    extractor = Agent(EXTRACT_SYSTEM, json_mode=True)\n",
    "    draft_date = json.loads(extractor.ask(INVOICES[name]))[\"due_date\"]\n",
    "    pred = json.loads(extractor.ask(\"Check your answer. Return the JSON again.\"))[\"due_date\"]\n",
    "    calls = calls + 2\n",
    "    ok = date_matches(pred, REFERENCES[name])\n",
    "    if ok:\n",
    "        correct = correct + 1\n",
    "    print(name, ok, REFERENCES[name], \"draft:\", draft_date, \"final:\", pred)\n",
    "print(\"accuracy:\", correct, \"/\", len(INVOICES), \"| model calls:\", calls)\n",
    "extractor.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29",
   "metadata": {},
   "source": [
    "Ten cases show how the measurement works; they do not establish performance on all invoices. Compare the `draft` and `final` columns. In our run, the check changed one date, and it changed a correct date into a wrong one (invoice I): the accuracy fell from 9 to 8 of 10, and the model calls doubled. The wrong date of invoice J stayed wrong. Compare with Part 1: there the agent searched between an evaluation and the next version, so new facts arrived. Here nothing new arrives. The extractor repeats the same date calculation, and the second calculation is not more reliable than the first. The evaluation set shows this; one example does not.\n",
    "\n",
    "### 2.4 ✍️ Grow the evaluation set\n",
    "\n",
    "*Do:* add one invoice of your own as `INVOICES[\"K\"]`, with its correct date in `REFERENCES[\"K\"]` (work it out yourself). Try a case you expect to be hard. Then rerun 2.2 and 2.3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ✍️ TYPE · exercise: write the code between the FILL IN lines yourself\n",
    "### FILL IN (START) ###\n",
    "# INVOICES[\"K\"] = \"Invoice K. ...\"\n",
    "# REFERENCES[\"K\"] = \"YYYY/MM/DD\"\n",
    "### FILL IN (END) ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · check: just run\n",
    "valid = True\n",
    "for name in REFERENCES:\n",
    "    try:\n",
    "        datetime.strptime(REFERENCES[name], \"%Y/%m/%d\")      # every reference must be a real date\n",
    "    except ValueError:\n",
    "        valid = False\n",
    "        print(\"not a valid reference:\", name, REFERENCES[name])\n",
    "done[\"2.4\"] = \"K\" in INVOICES and len(REFERENCES) == len(INVOICES) and valid\n",
    "done[\"2.4\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c32",
   "metadata": {},
   "source": [
    "## 3. An LLM judge\n",
    "\n",
    "Did the reflection of Part 1 make the essay better? The criteria for an essay need interpretation: \"the essay explains the cause step by step\" can be done in many different words, so a string comparison cannot check it. An **LLM judge** is a model asked to grade an output against written criteria; a **rubric** gives the criteria and the rule for the score. A judge is also a model, so it can be wrong. Before its scores count, compare them with a person's scores.\n",
    "\n",
    "### 3.1 The judge\n",
    "\n",
    "The rubric is the system message of the judge. Each essay arrives as a user message, and `json_mode=True` makes the reply valid JSON, so code can read the scores. Every essay gets a new judge. The writer in Part 1 never saw this rubric.\n",
    "\n",
    "*Do:* run the cell. It grades the essay of 1.2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "JUDGE_SYSTEM = \"\"\"You grade short science essays. Score each criterion from 1 (weak) to 5 (strong) and give a short reason.\n",
    "1. explanation: the essay answers the question and explains the cause step by step.\n",
    "2. sources: each main fact has a citation of a Wikipedia article, written as [article title].\n",
    "3. clarity: a general reader can follow the essay, and technical terms are explained.\n",
    "Return JSON: {\"explanation\": n, \"sources\": n, \"clarity\": n, \"reasons\": \"...\"}\"\"\"\n",
    "CRITERIA = [\"explanation\", \"sources\", \"clarity\"]\n",
    "\n",
    "def judge_essay(question, essay):\n",
    "    \"\"\"Score one essay with a new judge; return the JSON answer as a dictionary.\"\"\"\n",
    "    judge = Agent(JUDGE_SYSTEM, json_mode=True)    # a new judge for every essay\n",
    "    return json.loads(judge.ask(f\"\"\"Essay question: {question}\n",
    "\n",
    "Essay:\n",
    "{essay}\"\"\"))\n",
    "\n",
    "def scores_of(result):\n",
    "    scores = []\n",
    "    for name in CRITERIA:\n",
    "        scores.append(result[name])\n",
    "    return scores\n",
    "\n",
    "direct_result = judge_essay(QUESTION, direct_essay)\n",
    "print(scores_of(direct_result))\n",
    "print(direct_result[\"reasons\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c34",
   "metadata": {},
   "source": [
    "### 3.2 Check the judge against a person\n",
    "\n",
    "`LABELED` holds five short essays on the question of Part 1, each with the scores that a person gave on the same rubric. The judge **agrees** with the person on a criterion when the two scores differ by at most one point.\n",
    "\n",
    "*Do:* run the cell. Count the agreements and read each disagreement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "LABELED = [\n",
    "    (\"\"\"A black hole forms when a very massive star dies. During its life, the star burns fuel in its core, and the heat\n",
    "pushes outward against gravity [Stellar evolution]. When the fuel runs out, nothing holds the core up, and it collapses in\n",
    "less than a second [Supernova]. If the collapsed core is heavier than about three Suns, no known force can stop the\n",
    "collapse, and it becomes a black hole: an object whose gravity lets nothing escape, not even light [Black hole].\"\"\", [5, 5, 5]),\n",
    "    (\"\"\"A black hole forms when a very massive star dies. While the star burns fuel, the heat pushes outward and balances\n",
    "gravity. When the fuel runs out, the balance ends and the core collapses. If the core is heavy enough, nothing can stop\n",
    "the collapse, and it becomes a black hole, an object whose gravity lets nothing escape.\"\"\", [5, 1, 5]),\n",
    "    (\"\"\"A black hole is a region of space where gravity is so strong that nothing can escape from it, not even\n",
    "light [Black hole]. The edge of this region is called the event horizon [Event horizon]. The first image of a black hole\n",
    "was published in 2019 [Event Horizon Telescope].\"\"\", [1, 5, 5]),\n",
    "    (\"\"\"Black holes form via core collapse of progenitors above the Tolman-Oppenheimer-Volkoff limit once degeneracy\n",
    "pressure fails [Stellar black hole]. Post-main-sequence stars exhaust fusion, the iron core exceeds the Chandrasekhar\n",
    "mass, and collapse proceeds past the neutron-star stage to a singularity inside an event horizon [Black hole].\"\"\", [4, 5, 1]),\n",
    "    (\"\"\"Black holes are very mysterious. Scientists have studied them for many years, and there are many theories about\n",
    "them. They form in space in different ways, and they are important for our understanding of the universe.\"\"\", [1, 1, 4]),\n",
    "]\n",
    "\n",
    "agree = 0\n",
    "for essay_text, labels in LABELED:\n",
    "    scores = scores_of(judge_essay(QUESTION, essay_text))\n",
    "    for i in range(3):\n",
    "        if abs(scores[i] - labels[i]) <= 1:        # abs: the size of the difference\n",
    "            agree = agree + 1\n",
    "    print(\"judge:\", scores, \"| person:\", labels)\n",
    "print(\"agreement:\", agree, \"/\", 3 * len(LABELED))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c36",
   "metadata": {},
   "source": [
    "**Reading the output.** Most scores agree within one point. Read each disagreement. In our run, the two disagreements were on clarity: the judge gave 3 to the essay that is full of technical terms, where the person gave 1, and it gave 2 to the vague essay, which the person found easy to follow (4). A disagreement shows where the rubric needs clearer words, here what \"a general reader can follow\" means. Check the judge again when you change the rubric.\n",
    "\n",
    "### 3.3 Did the reflection help?\n",
    "\n",
    "The same judge now grades three texts on the same question: the essay without reflection (1.2), the first reply of the reflection agent (its draft, 1.3), and its final essay (1.4). The model calls show the cost: every assistant message in a list is one call.\n",
    "\n",
    "*Do:* run the cell. Compare the three lines of scores, then read the reasons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · example: read each line, then run\n",
    "def count_calls(agent):\n",
    "    \"\"\"Count the model calls of an agent: every assistant message is one call.\"\"\"\n",
    "    calls = 0\n",
    "    for m in agent.messages:\n",
    "        if m[\"role\"] == \"assistant\":\n",
    "            calls = calls + 1\n",
    "    return calls\n",
    "\n",
    "draft_result = judge_essay(QUESTION, draft)\n",
    "final_result = judge_essay(QUESTION, final_essay)\n",
    "print(\"without reflection (1.2):\", scores_of(direct_result), \"|\", len(direct_essay.split()), \"words |\", count_calls(writer), \"model calls\")\n",
    "print(\"first reply, draft (1.3):\", scores_of(draft_result), \"|\", len(draft.split()), \"words\")\n",
    "print(\"final essay (1.4):\", scores_of(final_result), \"|\", len(final_essay.split()), \"words |\", count_calls(agent), \"model calls\")\n",
    "print()\n",
    "print(\"WITHOUT REFLECTION:\", direct_result[\"reasons\"])\n",
    "print(\"DRAFT:\", draft_result[\"reasons\"])\n",
    "print(\"FINAL:\", final_result[\"reasons\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c38",
   "metadata": {},
   "source": [
    "**Reading the output.** In our runs, the first reply scored 1 or 2 on sources: the agent wrote that draft before it searched. The final essay scored 4 or 5 on sources: the evaluation named the missing citations, and the improved version added them. So the reflection repaired its own draft. Now compare with the first line. The writer without reflection searched before it wrote, and it also scored 4 or 5 on sources, with 2 to 4 model calls against about 10. On this question, the reflection gave one point more at most, and it cost several times the calls. Two cautions: 5 is the top of the scale, so this rubric cannot separate two good essays; and one question is not an evaluation set. Read the reasons, not only the numbers.\n",
    "\n",
    "### 3.4 ✍️ Judge your own essay\n",
    "\n",
    "*Do:* write a short essay on `QUESTION` in `MY_ESSAY` and your own three scores (1 to 5) in `MY_LABELS`. Then write the line that asks the judge, as in 3.2. Try an essay that might confuse the judge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ✍️ TYPE · exercise: write the code between the FILL IN lines yourself\n",
    "### FILL IN (START) ###\n",
    "MY_ESSAY = \"\"\n",
    "MY_LABELS = [1, 1, 1]\n",
    "### FILL IN (END) ###\n",
    "### FILL IN (START) ###\n",
    "my_scores = [0, 0, 0]                              # replace: the judge's scores for MY_ESSAY\n",
    "### FILL IN (END) ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · check: just run\n",
    "print(\"judge:\", my_scores, \"| you:\", MY_LABELS)\n",
    "done[\"3.4\"] = MY_ESSAY != \"\" and my_scores != [0, 0, 0]\n",
    "done[\"3.4\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c41",
   "metadata": {},
   "source": [
    "## 4. Explain in your own words\n",
    "\n",
    "Write two or three sentences for each in `answers` (not checked).\n",
    "\n",
    "1. In Part 1, what did the first evaluation ask for, and which of these changes are in the final essay?\n",
    "2. In Part 1, one object did every step. Where in its list do you see that it went back to writing after an evaluation, and what made it stop?\n",
    "3. Why did every invoice and every graded essay get a new object in Parts 2 and 3?\n",
    "4. In Part 3, what did you compare the judge with before you used its scores, and what did the reflection cost?\n",
    "\n",
    "*Do:* write your answers, then run the last cell to see which exercises passed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ✍️ WRITE · your answer in words (not checked)\n",
    "answers = \"\"\"\n",
    "1.\n",
    "\n",
    "2.\n",
    "\n",
    "3.\n",
    "\n",
    "4.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ▶ RUN · check: just run\n",
    "for key in done:\n",
    "    print(key, done[key])"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
