{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 7: Observability and Evaluation Stack\n",
    "**Goal:** make a RAG assistant *observable* (you can see what it did and how long it took) and *evaluated* (you can measure whether it is correct, grounded, fast and affordable), then use those numbers as a release gate.\n",
    "\n",
    "**Why:** an LLM system can fail silently. It does not crash; it gives a fluent wrong answer. Without traces and evals you find out from a customer. This is Module 8, sections 2 (traces, metrics, logs), 3 (evals), 8 (open-source tools such as Langfuse) and 10 (production dashboard).\n",
    "\n",
    "You will:\n",
    "1. build a small RAG assistant (a short recap of Lab 3),\n",
    "2. record a **trace** for every question,\n",
    "3. run a **golden set** of 20 questions,\n",
    "4. score each answer with an **LLM judge**,\n",
    "5. summarise latency, estimated cost and quality, and apply pass/fail **gates**,\n",
    "6. optionally send the same traces to **Langfuse**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 0: Choose your model provider\n",
    "**Why we do this:** the same provider setup as Labs 4 and 6. This lab makes about 40 model calls per run (20 answers and 20 judgements), so a free tier with tight limits may slow you down. You can lower `LIMIT` below to run fewer questions.\n",
    "\n",
    "**What the code does:** defines `chat(messages)`, then an `ask(system, user)` helper with a few retries and a short wait, because rate-limit errors are common when you run many calls in a row.\n",
    "\n",
    "Set `PROVIDER` first. Ollama must run locally (Colab cannot reach it). Model names change, so edit `MODEL` if one is rejected."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "%pip install -q chromadb pypdf openai anthropic"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os, json, getpass\n",
    "\n",
    "PROVIDER = \"groq\"   # ollama | groq | gemini | openrouter | openai | anthropic\n",
    "\n",
    "PROVIDERS = {\n",
    "    \"ollama\":     {\"base_url\": \"http://localhost:11434/v1\",                                \"key_env\": None,                \"model\": \"llama3.1:8b\"},\n",
    "    \"groq\":       {\"base_url\": \"https://api.groq.com/openai/v1\",                          \"key_env\": \"GROQ_API_KEY\",       \"model\": \"llama-3.3-70b-versatile\"},\n",
    "    \"gemini\":     {\"base_url\": \"https://generativelanguage.googleapis.com/v1beta/openai/\", \"key_env\": \"GEMINI_API_KEY\",     \"model\": \"gemini-2.5-flash\"},\n",
    "    \"openrouter\": {\"base_url\": \"https://openrouter.ai/api/v1\",                            \"key_env\": \"OPENROUTER_API_KEY\", \"model\": \"openrouter/free\"},\n",
    "    \"openai\":     {\"base_url\": None,                                                      \"key_env\": \"OPENAI_API_KEY\",     \"model\": \"gpt-4o-mini\"},\n",
    "    \"anthropic\":  {\"base_url\": None,                                                      \"key_env\": \"ANTHROPIC_API_KEY\",  \"model\": \"claude-sonnet-5-5\"},\n",
    "}\n",
    "cfg = PROVIDERS[PROVIDER]\n",
    "MODEL = cfg[\"model\"]\n",
    "if cfg[\"key_env\"] and not os.getenv(cfg[\"key_env\"]):\n",
    "    os.environ[cfg[\"key_env\"]] = getpass.getpass(cfg[\"key_env\"] + \": \")\n",
    "\n",
    "def to_anthropic(messages):\n",
    "    # Convert OpenAI-style messages into Anthropic's format.\n",
    "    system = messages[0][\"content\"]\n",
    "    out = []\n",
    "    for m in messages[1:]:\n",
    "        if m[\"role\"] == \"user\":\n",
    "            out.append({\"role\": \"user\", \"content\": m[\"content\"]})\n",
    "        elif m[\"role\"] == \"assistant\":\n",
    "            blocks = []\n",
    "            if m.get(\"content\"):\n",
    "                blocks.append({\"type\": \"text\", \"text\": m[\"content\"]})\n",
    "            for tc in m.get(\"tool_calls\", []):\n",
    "                blocks.append({\"type\": \"tool_use\", \"id\": tc[\"id\"], \"name\": tc[\"function\"][\"name\"],\n",
    "                               \"input\": json.loads(tc[\"function\"][\"arguments\"] or \"{}\")})\n",
    "            out.append({\"role\": \"assistant\", \"content\": blocks})\n",
    "        elif m[\"role\"] == \"tool\":\n",
    "            block = {\"type\": \"tool_result\", \"tool_use_id\": m[\"tool_call_id\"], \"content\": m[\"content\"]}\n",
    "            if out and out[-1][\"role\"] == \"user\" and isinstance(out[-1][\"content\"], list):\n",
    "                out[-1][\"content\"].append(block)\n",
    "            else:\n",
    "                out.append({\"role\": \"user\", \"content\": [block]})\n",
    "    return system, out\n",
    "\n",
    "def chat(messages, tools=None):\n",
    "    # Returns {\"role\": \"assistant\", \"content\": str, \"tool_calls\": [...] (optional)}\n",
    "    if PROVIDER == \"anthropic\":\n",
    "        import anthropic\n",
    "        system, msgs = to_anthropic(messages)\n",
    "        kw = dict(model=MODEL, max_tokens=1000, temperature=0, system=system, messages=msgs)\n",
    "        if tools:\n",
    "            kw[\"tools\"] = [{\"name\": t[\"function\"][\"name\"], \"description\": t[\"function\"][\"description\"],\n",
    "                            \"input_schema\": t[\"function\"][\"parameters\"]} for t in tools]\n",
    "        r = anthropic.Anthropic().messages.create(**kw)\n",
    "        msg = {\"role\": \"assistant\", \"content\": \"\".join(b.text for b in r.content if b.type == \"text\")}\n",
    "        calls = [{\"id\": b.id, \"type\": \"function\", \"function\": {\"name\": b.name, \"arguments\": json.dumps(b.input)}}\n",
    "                 for b in r.content if b.type == \"tool_use\"]\n",
    "        if calls:\n",
    "            msg[\"tool_calls\"] = calls\n",
    "        return msg\n",
    "    from openai import OpenAI\n",
    "    key = os.getenv(cfg[\"key_env\"]) if cfg[\"key_env\"] else \"ollama\"\n",
    "    client = OpenAI(base_url=cfg[\"base_url\"], api_key=key) if cfg[\"base_url\"] else OpenAI(api_key=key)\n",
    "    kw = dict(model=MODEL, temperature=0, messages=messages)\n",
    "    if tools:\n",
    "        kw[\"tools\"] = tools\n",
    "    m = client.chat.completions.create(**kw).choices[0].message\n",
    "    msg = {\"role\": \"assistant\", \"content\": m.content or \"\"}\n",
    "    if m.tool_calls:\n",
    "        msg[\"tool_calls\"] = [{\"id\": tc.id, \"type\": \"function\",\n",
    "                              \"function\": {\"name\": tc.function.name, \"arguments\": tc.function.arguments}}\n",
    "                             for tc in m.tool_calls]\n",
    "    return msg\n",
    "\n",
    "print(chat([{\"role\": \"system\", \"content\": \"Be brief.\"}, {\"role\": \"user\", \"content\": \"Say hello in five words.\"}])[\"content\"])"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def ask(system: str, user: str, retries: int = 3) -> str:\n",
    "    for attempt in range(retries):\n",
    "        try:\n",
    "            return chat([{\"role\": \"system\", \"content\": system}, {\"role\": \"user\", \"content\": user}])[\"content\"]\n",
    "        except Exception as e:\n",
    "            if attempt == retries - 1:\n",
    "                raise\n",
    "            time.sleep(2 ** (attempt + 1))   # back off: 2s, 4s\n",
    "\n",
    "LIMIT = 20    # how many golden questions to run (lower it to save calls)\n",
    "PAUSE = 0.5   # seconds to wait between questions (raise it on strict free tiers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1: Build the assistant (recap of Lab 3)\n",
    "**Why we do this:** you need something to observe. This is a compact version of Lab 3: download the SEBI memo, chunk it, store it in ChromaDB, retrieve, and answer from the retrieved text only.\n",
    "\n",
    "**What the code does:** `retrieve_chunks` returns the top matching chunks with page numbers. `generate` builds a prompt from them and asks the model, replying \"Not found in the provided documents.\" when the context does not contain the answer. It returns both the answer and the prompt, because the next exercise needs the prompt to estimate tokens.\n",
    "\n",
    "**Key terms:** *Token* = the unit models read and are billed in, roughly three-quarters of an English word."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import re, urllib.request\n",
    "import chromadb\n",
    "from pypdf import PdfReader\n",
    "\n",
    "os.makedirs(\"data\", exist_ok=True)\n",
    "SEBI_URL = \"https://www.sebi.gov.in/sebi_data/meetingfiles/dec-2024/1735042007618_1.pdf\"\n",
    "if not os.path.exists(\"data/sebi_ai_memo.pdf\"):\n",
    "    req = urllib.request.Request(SEBI_URL, headers={\"User-Agent\": \"Mozilla/5.0\"})\n",
    "    open(\"data/sebi_ai_memo.pdf\", \"wb\").write(urllib.request.urlopen(req).read())\n",
    "\n",
    "PAGES = {i + 1: re.sub(r\"\\s+\", \" \", p.extract_text() or \"\").strip() for i, p in enumerate(PdfReader(\"data/sebi_ai_memo.pdf\").pages)}\n",
    "\n",
    "def chunk(text, size=600, overlap=100):\n",
    "    return [text[s:s + size] for s in range(0, max(len(text), 1), size - overlap) if text[s:s + size].strip()]\n",
    "\n",
    "col = chromadb.Client().get_or_create_collection(\"lab7\")\n",
    "ids, docs, metas = [], [], []\n",
    "for pg, text in PAGES.items():\n",
    "    for j, piece in enumerate(chunk(text)):\n",
    "        ids.append(f\"p{pg}-{j}\"); docs.append(piece); metas.append({\"page\": pg})\n",
    "col.add(ids=ids, documents=docs, metadatas=metas)\n",
    "\n",
    "def retrieve_chunks(question, k=4):\n",
    "    r = col.query(query_texts=[question], n_results=k)\n",
    "    return [{\"text\": t, \"page\": m[\"page\"]} for t, m in zip(r[\"documents\"][0], r[\"metadatas\"][0])]\n",
    "\n",
    "SYSTEM = (\"You answer questions about the SEBI memo using ONLY the numbered sources. Cite pages like (p.3). \"\n",
    "          \"If the sources do not contain the answer, reply exactly: 'Not found in the provided documents.'\")\n",
    "\n",
    "def generate(question, chunks):\n",
    "    sources = \"\\n\\n\".join(f\"[{i+1}] (p.{c['page']}) {c['text']}\" for i, c in enumerate(chunks))\n",
    "    prompt = f\"Sources:\\n{sources}\\n\\nQuestion: {question}\"\n",
    "    return {\"answer\": ask(SYSTEM, prompt), \"prompt\": SYSTEM + \"\\n\" + prompt}\n",
    "\n",
    "q = \"In which year did SEBI circulars mandate reporting of AI and ML applications?\"\n",
    "g = generate(q, retrieve_chunks(q))\n",
    "print(g[\"answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** an answer mentioning 2019, with a page citation. If it says \"Not found\", check that the chunk count is above zero and try a stronger model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2: Add tracing\n",
    "**Why we do this:** when an answer is wrong or slow you need to know *where*: did retrieval fail, or did generation? A **trace** records one request from start to finish. Inside it, each step is a **span** with its own timing and details (Module 8, section 2).\n",
    "\n",
    "**What the code does:** the `Trace` class below records spans. Its `span(...)` context manager times whatever runs inside the `with` block and stores the result. `finish` adds up time, estimated tokens and estimated cost for the whole request.\n",
    "\n",
    "**Key terms:**\n",
    "- *Trace:* the full record of one request.\n",
    "- *Span:* one timed step inside a trace.\n",
    "- *Estimated tokens:* we use about 4 characters per token. These are **estimates, not billed numbers**. Your provider's usage fields are exact (a stretch goal).\n",
    "\n",
    "**Your task:** complete `rag_answer` so it records a `retrieve` span and a `generate` span. In the generate span, store `in_tok` and `out_tok` in the span dictionary.\n",
    "\n",
    "**Cost:** set `PRICE_IN_PER_M` and `PRICE_OUT_PER_M` (dollars per million tokens) from your provider's current pricing page. Leave them at 0 for free tiers."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import uuid, contextlib, math\n",
    "\n",
    "PRICE_IN_PER_M = 0.0    # dollars per million input tokens (set from your provider's pricing page)\n",
    "PRICE_OUT_PER_M = 0.0   # dollars per million output tokens\n",
    "\n",
    "def est_tokens(text: str) -> int:\n",
    "    return max(1, len(text) // 4)\n",
    "\n",
    "TRACES = []\n",
    "\n",
    "class Trace:\n",
    "    def __init__(self, name, question):\n",
    "        self.id = uuid.uuid4().hex[:8]\n",
    "        self.name, self.question = name, question\n",
    "        self.spans, self.output = [], None\n",
    "        self.total_ms = self.in_tok = self.out_tok = 0\n",
    "        self.cost = 0.0\n",
    "\n",
    "    @contextlib.contextmanager\n",
    "    def span(self, name, **attrs):\n",
    "        s = {\"name\": name, **attrs}\n",
    "        t0 = time.time()\n",
    "        try:\n",
    "            yield s\n",
    "        finally:\n",
    "            s[\"ms\"] = round((time.time() - t0) * 1000)\n",
    "            self.spans.append(s)\n",
    "\n",
    "    def finish(self, output):\n",
    "        self.output = output\n",
    "        self.total_ms = sum(s[\"ms\"] for s in self.spans)\n",
    "        self.in_tok = sum(s.get(\"in_tok\", 0) for s in self.spans)\n",
    "        self.out_tok = sum(s.get(\"out_tok\", 0) for s in self.spans)\n",
    "        self.cost = (self.in_tok * PRICE_IN_PER_M + self.out_tok * PRICE_OUT_PER_M) / 1_000_000\n",
    "        TRACES.append(self)\n",
    "\n",
    "def rag_answer(question, k=4):\n",
    "    trace = Trace(\"rag_answer\", question)\n",
    "    # YOUR CODE HERE:\n",
    "    # 1. with trace.span(\"retrieve\", k=k) as s: retrieve the chunks, then set s[\"n_chunks\"]\n",
    "    # 2. with trace.span(\"generate\", model=MODEL) as s: call generate(question, chunks),\n",
    "    #    then set s[\"in_tok\"] and s[\"out_tok\"] using est_tokens on the prompt and the answer\n",
    "    # 3. call trace.finish(answer) and return {\"answer\": ..., \"chunks\": ..., \"trace\": trace}\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details><summary>Solution (try first)</summary>\n",
    "\n",
    "```python\n",
    "def rag_answer(question, k=4):\n",
    "    trace = Trace(\"rag_answer\", question)\n",
    "    with trace.span(\"retrieve\", k=k) as s:\n",
    "        chunks = retrieve_chunks(question, k)\n",
    "        s[\"n_chunks\"] = len(chunks)\n",
    "    with trace.span(\"generate\", model=MODEL) as s:\n",
    "        g = generate(question, chunks)\n",
    "        s[\"in_tok\"] = est_tokens(g[\"prompt\"])\n",
    "        s[\"out_tok\"] = est_tokens(g[\"answer\"])\n",
    "    trace.finish(g[\"answer\"])\n",
    "    return {\"answer\": g[\"answer\"], \"chunks\": chunks, \"trace\": trace}\n",
    "```\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "r = rag_answer(\"Which circular covered mutual funds and what was its date?\")\n",
    "t = r[\"trace\"]\n",
    "print(r[\"answer\"], \"\\n\")\n",
    "print(\"trace\", t.id, \"| total\", t.total_ms, \"ms | est tokens in/out:\", t.in_tok, \"/\", t.out_tok)\n",
    "for s in t.spans:\n",
    "    print(\"  span:\", s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** the answer, then one line per span showing its name, timing in milliseconds, and details. The `generate` span will almost always be far slower than `retrieve`.\n",
    "\n",
    "**What to notice:** that gap tells you where optimisation effort would pay off. If `retrieve` were the slow one, you would tune the database, not the model.\n",
    "\n",
    "**Finance link:** a trace is also an audit record: which question, which sources, which model, when (Module 8, section 9)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3: The golden set\n",
    "**Why we do this:** you cannot measure quality without a fixed set of questions with known right answers. This is a *golden set*. Running the same set after every change shows whether you made things better or worse (a *regression test*).\n",
    "\n",
    "**What the code does:** 20 questions written from the SEBI memo you loaded. 18 have answers in the document. 2 do not (the memo excises the number of comments received, and states no penalty). For those, the *correct* behaviour is to say \"Not found\". Testing that is important: a system that always answers is dangerous in finance.\n",
    "\n",
    "Each item has `keys`: a list of groups. An answer passes if, for every group, at least one alternative appears in it (case-insensitive). This is a simple, repeatable fact check.\n",
    "\n",
    "**Your task:** add two questions of your own at the end of the list (with `keys`), taken from the memo. Check the memo text first so the expected answer is right.\n",
    "\n",
    "**Important:** a golden set of 20 is enough to teach the idea and to catch big regressions, but it is far too small to detect small quality changes. Real teams grow it from production failures."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "GOLDEN = [\n",
    "    {\"q\": \"Which three SEBI regulations does the memo propose to amend?\", \"answerable\": True,\n",
    "     \"keys\": [[\"intermediaries\"], [\"stock exchanges\"], [\"depositories\"]]},\n",
    "    {\"q\": \"What is the date of the SEBI consultation paper on this proposal?\", \"answerable\": True,\n",
    "     \"keys\": [[\"november 13, 2024\", \"13 november 2024\", \"13th november 2024\", \"nov 13, 2024\", \"13/11/2024\"]]},\n",
    "    {\"q\": \"In which year did SEBI circulars first mandate reporting of AI and ML applications?\", \"answerable\": True,\n",
    "     \"keys\": [[\"2019\"]]},\n",
    "    {\"q\": \"What is the date of the circular covering stock brokers and depository participants?\", \"answerable\": True,\n",
    "     \"keys\": [[\"january 4, 2019\", \"4 january 2019\", \"4th january 2019\"]]},\n",
    "    {\"q\": \"What is the date of the circular covering market infrastructure institutions?\", \"answerable\": True,\n",
    "     \"keys\": [[\"january 31, 2019\", \"31 january 2019\", \"january, 31, 2019\"]]},\n",
    "    {\"q\": \"What is the date of the circular covering mutual funds and asset management companies?\", \"answerable\": True,\n",
    "     \"keys\": [[\"may 9, 2019\", \"9 may 2019\"]]},\n",
    "    {\"q\": \"For which purposes do stock brokers and depository participants use AI/ML tools?\", \"answerable\": True,\n",
    "     \"keys\": [[\"surveillance\"], [\"order execution\"], [\"kyc\"]]},\n",
    "    {\"q\": \"For which purposes do mutual funds use AI/ML tools?\", \"answerable\": True,\n",
    "     \"keys\": [[\"surveillance\"], [\"customer service\"], [\"customer segmentation\"]]},\n",
    "    {\"q\": \"Which kinds of institutions are market infrastructure institutions?\", \"answerable\": True,\n",
    "     \"keys\": [[\"stock exchange\"], [\"depositor\"], [\"clearing corporation\"]]},\n",
    "    {\"q\": \"What three things would a regulated person be solely responsible for under the proposal?\", \"answerable\": True,\n",
    "     \"keys\": [[\"privacy\", \"data\"], [\"output\"], [\"compliance\", \"applicable laws\"]]},\n",
    "    {\"q\": \"Does the proposed responsibility depend on the scale of AI adoption?\", \"answerable\": True,\n",
    "     \"keys\": [[\"irrespective\", \"regardless\", \"does not depend\", \"not depend\", \"any scale\", \"scale and scenario\"]]},\n",
    "    {\"q\": \"Does the proposal cover AI tools bought from third-party providers?\", \"answerable\": True,\n",
    "     \"keys\": [[\"third-party\", \"third party\"]]},\n",
    "    {\"q\": \"What does the memo ask the Board to approve?\", \"answerable\": True,\n",
    "     \"keys\": [[\"amend\"], [\"chairperson\"]]},\n",
    "    {\"q\": \"Why was the AI/ML reporting mechanism introduced?\", \"answerable\": True,\n",
    "     \"keys\": [[\"adoption\"], [\"preparedness\", \"prepared\"]]},\n",
    "    {\"q\": \"Name some benefits of AI that the memo lists.\", \"answerable\": True,\n",
    "     \"keys\": [[\"efficiency\"], [\"cyber security\", \"cybersecurity\"], [\"customer experience\"]]},\n",
    "    {\"q\": \"Which AI use by market infrastructure institutions involves generative AI?\", \"answerable\": True,\n",
    "     \"keys\": [[\"generative\"], [\"document\", \"report\", \"order\"]]},\n",
    "    {\"q\": \"Where is the consultation paper available?\", \"answerable\": True,\n",
    "     \"keys\": [[\"annexure a\", \"reports for public comments\", \"website\", \"sebi.gov.in\"]]},\n",
    "    {\"q\": \"Why does SEBI want to assign responsibility to regulated persons using AI?\", \"answerable\": True,\n",
    "     \"keys\": [[\"investor\"], [\"serious\"]]},\n",
    "    {\"q\": \"What penalty will apply to a regulated person who breaches the proposed rule?\", \"answerable\": False, \"keys\": []},\n",
    "    {\"q\": \"How many public comments were received on the consultation paper?\", \"answerable\": False, \"keys\": []},\n",
    "    # YOUR CODE HERE: add two questions of your own from the memo, in the same format\n",
    "]\n",
    "\n",
    "def check_facts(answer: str, keys) -> bool:\n",
    "    a = answer.lower()\n",
    "    return all(any(alt in a for alt in group) for group in keys)\n",
    "\n",
    "print(len(GOLDEN), \"golden questions,\", sum(g[\"answerable\"] for g in GOLDEN), \"answerable\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** `20 golden questions, 18 answerable` (more if you added your own).\n",
    "\n",
    "**What to notice:** the fact check is strict about wording. A correct answer phrased in an unexpected way would fail it. That is why we also use an LLM judge, and why some failures deserve a human look."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 4: Evaluate with an LLM judge\n",
    "**Why we do this:** the fact check tells you if the key facts are present. It does not tell you if the answer *invented* anything else. **Groundedness** asks: is every claim in the answer supported by the retrieved text? A model can be correct yet partly invented, and in finance that matters (Module 8, section 3).\n",
    "\n",
    "**What the code does:** an *LLM judge* is a second model call that reads the sources and the answer and gives a 1 (fully supported) or 0 (some unsupported claim). `evaluate` runs every golden question through `rag_answer`, then records correctness, groundedness, latency and cost.\n",
    "\n",
    "**Key terms:**\n",
    "- *LLM-as-judge:* using a model to grade another model's output.\n",
    "- *Groundedness (faithfulness):* the answer sticks to the supplied sources.\n",
    "\n",
    "**Your task:** complete `groundedness(context, answer)`. Write a judge prompt that ends with `SCORE: 1` or `SCORE: 0` followed by `REASON:`, call `ask`, and parse the score with a regular expression. Return 1 or 0, and return 0 if no score can be parsed.\n",
    "\n",
    "**Important:** judges are imperfect. They can be lenient with fluent text, and if the judge is the same model as the generator it tends to favour its own style. Use a different, stronger model as judge when you can. Spot-check a few judgements by hand."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def groundedness(context: str, answer: str) -> int:\n",
    "    # YOUR CODE HERE:\n",
    "    # 1. build a system prompt telling the judge to check ONLY whether every factual claim\n",
    "    #    in the answer is supported by the context (an answer saying 'Not found' is grounded)\n",
    "    # 2. ask it to reply with 'SCORE: 1' or 'SCORE: 0' then 'REASON: <one sentence>'\n",
    "    # 3. parse with re.search(r\"SCORE:\\s*([01])\", reply); return int(...) or 0 if no match\n",
    "    raise NotImplementedError\n",
    "\n",
    "def percentile(values, p):\n",
    "    s = sorted(values)\n",
    "    return s[max(0, math.ceil(p / 100 * len(s)) - 1)]\n",
    "\n",
    "def evaluate(items, k=4, label=\"run\"):\n",
    "    rows = []\n",
    "    for it in items[:LIMIT]:\n",
    "        r = rag_answer(it[\"q\"], k)\n",
    "        ans, t = r[\"answer\"], r[\"trace\"]\n",
    "        correct = check_facts(ans, it[\"keys\"]) if it[\"answerable\"] else (\"not found\" in ans.lower())\n",
    "        context = \"\\n\".join(ch[\"text\"] for ch in r[\"chunks\"])\n",
    "        rows.append({\"q\": it[\"q\"], \"answer\": ans, \"correct\": int(correct),\n",
    "                     \"grounded\": groundedness(context, ans), \"ms\": t.total_ms, \"cost\": t.cost})\n",
    "        time.sleep(PAUSE)\n",
    "    print(f\"[{label}] evaluated {len(rows)} questions\")\n",
    "    return rows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details><summary>Solution (try first)</summary>\n",
    "\n",
    "```python\n",
    "def groundedness(context, answer):\n",
    "    system = (\"You are a strict fact-checker. Decide ONLY whether every factual claim in the ANSWER \"\n",
    "              \"is supported by the CONTEXT. An answer saying 'Not found in the provided documents.' is grounded. \"\n",
    "              \"Reply with 'SCORE: 1' if all claims are supported or 'SCORE: 0' if any is not, \"\n",
    "              \"then 'REASON:' and one sentence.\")\n",
    "    reply = ask(system, f\"CONTEXT:\\n{context}\\n\\nANSWER:\\n{answer}\")\n",
    "    m = re.search(r\"SCORE:\\s*([01])\", reply)\n",
    "    return int(m.group(1)) if m else 0\n",
    "```\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "rows = evaluate(GOLDEN, k=4, label=\"baseline k=4\")\n",
    "for r in rows:\n",
    "    flag = \"ok \" if r[\"correct\"] and r[\"grounded\"] else \"FAIL\"\n",
    "    print(f\"{flag}  correct={r['correct']} grounded={r['grounded']} {r['ms']:>5}ms  {r['q'][:70]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** one line per question. Most should pass; a few `FAIL` lines are normal and are the point.\n",
    "\n",
    "**What to notice:** read two or three failing answers in full (`rows[i][\"answer\"]`). Decide for yourself: is it a real error, a retrieval miss, a too-strict fact check, or a judge mistake? Those four causes call for four different fixes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 5: Summary, gates and a regression test\n",
    "**Why we do this:** 20 rows are hard to act on. A production dashboard reduces them to a few numbers with thresholds, so a person (or a CI job) can say \"ship\" or \"do not ship\" (Module 8, section 10).\n",
    "\n",
    "**What the code does:** you write `summarise(rows)` to compute correctness rate, groundedness rate, median and 95th-percentile latency, and cost per query. `gates` then compares them with thresholds. **These thresholds are illustrative.** Set yours from your own risk appetite and review them with your risk team.\n",
    "\n",
    "**Key terms:**\n",
    "- *p50 / p95 latency:* half / 95% of requests finish faster than this. p95 shows the slow tail that users actually notice.\n",
    "- *Gate:* a pass/fail threshold that blocks a release.\n",
    "- *Regression test:* re-running the golden set after a change to see if anything got worse.\n",
    "\n",
    "**Your task:** complete `summarise`. Hint: use `sum(...) / len(rows)` for the rates, and the `percentile` helper for latency."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def summarise(rows):\n",
    "    # YOUR CODE HERE: return a dict with keys\n",
    "    # correctness, groundedness, p50_ms, p95_ms, cost_per_query\n",
    "    raise NotImplementedError\n",
    "\n",
    "THRESHOLDS = {\"correctness\": 0.80, \"groundedness\": 0.90, \"p95_ms\": 8000, \"cost_per_query\": 0.01}   # illustrative\n",
    "\n",
    "def gates(s):\n",
    "    checks = {\n",
    "        \"correctness >= 0.80\":     s[\"correctness\"] >= THRESHOLDS[\"correctness\"],\n",
    "        \"groundedness >= 0.90\":    s[\"groundedness\"] >= THRESHOLDS[\"groundedness\"],\n",
    "        \"p95 latency <= 8000 ms\":  s[\"p95_ms\"] <= THRESHOLDS[\"p95_ms\"],\n",
    "        \"cost/query <= $0.01\":     s[\"cost_per_query\"] <= THRESHOLDS[\"cost_per_query\"],\n",
    "    }\n",
    "    for name, ok in checks.items():\n",
    "        print((\"PASS  \" if ok else \"FAIL  \") + name)\n",
    "    print(\"RELEASE:\", \"OK\" if all(checks.values()) else \"BLOCKED\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details><summary>Solution (try first)</summary>\n",
    "\n",
    "```python\n",
    "def summarise(rows):\n",
    "    n = len(rows)\n",
    "    return {\n",
    "        \"correctness\": sum(r[\"correct\"] for r in rows) / n,\n",
    "        \"groundedness\": sum(r[\"grounded\"] for r in rows) / n,\n",
    "        \"p50_ms\": percentile([r[\"ms\"] for r in rows], 50),\n",
    "        \"p95_ms\": percentile([r[\"ms\"] for r in rows], 95),\n",
    "        \"cost_per_query\": sum(r[\"cost\"] for r in rows) / n,\n",
    "    }\n",
    "```\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "base = summarise(rows)\n",
    "print({k: round(v, 4) for k, v in base.items()})\n",
    "gates(base)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Now run a regression test.** Change one setting, here retrieving only 1 chunk instead of 4, and compare. This simulates a \"cost-saving\" change that someone might ship without testing."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "rows_k1 = evaluate(GOLDEN, k=1, label=\"candidate k=1\")\n",
    "cand = summarise(rows_k1)\n",
    "print(f\"{'metric':16}{'baseline':>12}{'candidate':>12}\")\n",
    "for k in base:\n",
    "    print(f\"{k:16}{base[k]:>12.3f}{cand[k]:>12.3f}\")\n",
    "gates(cand)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** a side-by-side table. With only one chunk, correctness usually drops (questions needing several pieces of the memo suffer), while latency and cost may improve slightly.\n",
    "\n",
    "**What to notice:** without the golden set you would see only the lower cost and think the change was free. With only 20 questions, treat small differences as noise and big drops as real.\n",
    "\n",
    "**Finance link:** this before-and-after table is the kind of validation evidence that model-risk reviewers expect for any change to a model-driven process (Module 5, section 2; Module 8, section 9)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 6 (optional): Send traces to Langfuse\n",
    "**Why we do this:** your `Trace` class teaches what a trace is. In practice you use a tool that stores traces, draws them as a timeline, and lets you search and compare them. **Langfuse** is an open-source option (Module 8, section 8).\n",
    "\n",
    "**Two ways to run it:**\n",
    "- **Langfuse Cloud (works on Colab).** Create a free account and project at langfuse.com, then copy the public key, secret key and host URL. **Warning:** traces include your questions and answers and are sent to a third party. Never send real client data. Use only fake or public data like this lab's.\n",
    "- **Self-hosted with Docker (keeps data on your machine).** Docker must be installed. Clone the Langfuse repository and start it with `docker compose up`, following the current self-hosting guide on the Langfuse site. It runs several services and needs a few GB of memory, and it cannot be reached from Colab, so use a local notebook. The UI is normally at http://localhost:3000, where you create a project and keys.\n",
    "\n",
    "**What the code does:** reads your keys, then wraps the existing `retrieve_chunks` and `generate` functions with Langfuse's `observe` decorator. Because `rag_answer` calls them, each question now appears as a trace with nested `retrieve` and `generate` steps.\n",
    "\n",
    "**Important:** the Langfuse SDK changes between versions. This cell was written for the v3 Python SDK (`from langfuse import observe, get_client`). If an import or call fails, check the current Langfuse documentation. The rest of the lab does not depend on this cell."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "%pip install -q langfuse"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "LANGFUSE_PUBLIC_KEY = \"\"   # paste from your Langfuse project (leave empty to skip this exercise)\n",
    "LANGFUSE_SECRET_KEY = \"\"\n",
    "LANGFUSE_HOST = \"https://cloud.langfuse.com\"   # or your region's URL, or http://localhost:3000 when self-hosting\n",
    "\n",
    "if LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY:\n",
    "    os.environ[\"LANGFUSE_PUBLIC_KEY\"] = LANGFUSE_PUBLIC_KEY\n",
    "    os.environ[\"LANGFUSE_SECRET_KEY\"] = LANGFUSE_SECRET_KEY\n",
    "    os.environ[\"LANGFUSE_HOST\"] = LANGFUSE_HOST\n",
    "    try:\n",
    "        from langfuse import observe, get_client\n",
    "        retrieve_chunks = observe(name=\"retrieve\")(retrieve_chunks)\n",
    "        generate = observe(name=\"generate\")(generate)\n",
    "        for it in GOLDEN[:3]:\n",
    "            rag_answer(it[\"q\"])\n",
    "        get_client().flush()\n",
    "        print(\"Sent 3 traces. Open your Langfuse project and look at the Traces page.\")\n",
    "    except Exception as e:\n",
    "        print(\"Langfuse step skipped:\", type(e).__name__, e)\n",
    "else:\n",
    "    print(\"No Langfuse keys set. Skipping this optional exercise.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** \"Sent 3 traces\", and the Traces page in Langfuse showing three entries, each with its timing.\n",
    "\n",
    "**What to notice:** compare Langfuse's timeline with the spans your own `Trace` recorded. Your class already stores the same facts. The tool adds storage, search, dashboards and team sharing.\n",
    "\n",
    "**Finance link:** question text often contains personal data. Before sending traces anywhere, redact them (reuse `redact` from Lab 6) and check your retention and residency rules (Module 8, section 9).\n",
    "\n",
    "## Stretch goals\n",
    "1. Use your provider's exact token counts instead of the 4-characters-per-token estimate, and compare the two.\n",
    "2. Use a different, stronger model as the judge and see how many judgements change.\n",
    "3. Try DeepEval or RAGAS, open-source evaluation libraries, to compute faithfulness and answer relevance, and compare with your own judge.\n",
    "4. Redact personal data from traces before logging them.\n",
    "5. Add a latency or groundedness check that fails a script with a non-zero exit code, so a CI pipeline can block a release.\n",
    "\n",
    "## Common errors\n",
    "- **`NotImplementedError`:** you have not completed `rag_answer`, `groundedness` or `summarise` yet. Complete them in order.\n",
    "- **Rate limit or 429 errors:** raise `PAUSE`, lower `LIMIT`, or use a provider with higher limits. Each question uses two model calls.\n",
    "- **Groundedness is 0 for everything:** your judge reply is probably not in the `SCORE: 1` format. Print one raw reply and adjust the prompt or the regular expression.\n",
    "- **Langfuse import errors:** the SDK API changes between versions. Check the current docs, or skip Exercise 6.\n",
    "- **Scores change between runs:** expected. Model output varies, so repeat the run before drawing conclusions."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}