{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 3: Regulatory Document RAG\n",
    "**Goal:** answer questions about real financial documents, with citations, using ChromaDB and an LLM of your choice.\n",
    "\n",
    "**Why:** an LLM alone can invent regulatory details. RAG (Retrieval-Augmented Generation) first *finds* the relevant passages in your documents, then asks the model to answer only from them. This is Module 3, sections 4-6."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 0: Choose your provider\n",
    "**Why we do this:** the rest of the lab calls one function, `ask_llm()`. You pick the provider once here and nothing else changes.\n",
    "\n",
    "**What the code does:** it builds `ask_llm(system, user)` for Ollama, Groq or Anthropic. Ollama and Groq both speak the OpenAI-style API, so they share code.\n",
    "\n",
    "**Important:** Ollama runs on *your* computer. Google Colab cannot reach it. If you pick Ollama, run this notebook locally (Jupyter/VS Code). Use Groq or Anthropic on Colab.\n",
    "\n",
    "Model names change over time. If one fails, check the provider's model list and update `MODEL`."
   ]
  },
  {
   "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, getpass\n",
    "\n",
    "PROVIDER = \"groq\"   # \"ollama\" | \"groq\" | \"anthropic\"\n",
    "\n",
    "DEFAULTS = {\n",
    "    \"ollama\":    {\"model\": \"llama3.1:8b\",              \"base_url\": \"http://localhost:11434/v1\"},\n",
    "    \"groq\":      {\"model\": \"llama-3.3-70b-versatile\",  \"base_url\": \"https://api.groq.com/openai/v1\"},\n",
    "    \"anthropic\": {\"model\": \"claude-sonnet-5-5\",        \"base_url\": None},\n",
    "}\n",
    "MODEL = DEFAULTS[PROVIDER][\"model\"]\n",
    "\n",
    "if PROVIDER == \"groq\" and not os.getenv(\"GROQ_API_KEY\"):\n",
    "    os.environ[\"GROQ_API_KEY\"] = getpass.getpass(\"Groq API key: \")\n",
    "if PROVIDER == \"anthropic\" and not os.getenv(\"ANTHROPIC_API_KEY\"):\n",
    "    os.environ[\"ANTHROPIC_API_KEY\"] = getpass.getpass(\"Anthropic API key: \")\n",
    "\n",
    "def ask_llm(system: str, user: str) -> str:\n",
    "    if PROVIDER == \"anthropic\":\n",
    "        import anthropic\n",
    "        r = anthropic.Anthropic().messages.create(\n",
    "            model=MODEL, max_tokens=800, temperature=0, system=system,\n",
    "            messages=[{\"role\": \"user\", \"content\": user}])\n",
    "        return r.content[0].text\n",
    "    from openai import OpenAI\n",
    "    key = os.getenv(\"GROQ_API_KEY\") if PROVIDER == \"groq\" else \"ollama\"\n",
    "    client = OpenAI(base_url=DEFAULTS[PROVIDER][\"base_url\"], api_key=key)\n",
    "    r = client.chat.completions.create(\n",
    "        model=MODEL, temperature=0,\n",
    "        messages=[{\"role\": \"system\", \"content\": system}, {\"role\": \"user\", \"content\": user}])\n",
    "    return r.choices[0].message.content\n",
    "\n",
    "print(ask_llm(\"Be brief.\", \"Say hello in five words.\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** a short greeting. If you see an authentication or connection error, re-check your key (Groq/Anthropic) or that `ollama serve` is running and the model is pulled (Ollama)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1: Load the documents\n",
    "**Why we do this:** RAG quality starts with clean text. We keep the *page number* with every piece of text so answers can cite it later.\n",
    "\n",
    "**What the code does:** `load_pdf` reads each page with `pypdf` and returns a list of `(page_number, text)`.\n",
    "\n",
    "**Key terms:** *Corpus* = your collection of documents.\n",
    "\n",
    "**Documents:**\n",
    "- SEBI Board memo on responsibility for AI tools (downloaded for you below). It is a proposal paper, not a final circular.\n",
    "- RBI Master Direction on IT Governance (2023): download from https://www.rbi.org.in/Scripts/BS_ViewMasDirections.aspx?id=12562\n",
    "- HDFC Bank Integrated Annual Report 2024-25 (large; optional): https://www.hdfcbank.com/personal/about-us/investor-relations/annual-reports\n",
    "\n",
    "Put the PDFs in a `data/` folder next to this notebook (on Colab, use the Files panel to upload). RBI's site blocks scripted downloads, so download that file in your browser."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os, urllib.request\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",
    "from pypdf import PdfReader\n",
    "def load_pdf(path):\n",
    "    return [(i + 1, (p.extract_text() or \"\")) for i, p in enumerate(PdfReader(path).pages)]\n",
    "\n",
    "DOCS = {f: load_pdf(f\"data/{f}\") for f in os.listdir(\"data\") if f.lower().endswith(\".pdf\")}\n",
    "for name, pages in DOCS.items():\n",
    "    print(name, \"-\", len(pages), \"pages,\", sum(len(t) for _, t in pages), \"characters\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** one line per PDF with its page and character counts. The SEBI memo has 6 pages. If a PDF shows near-zero characters it is a scanned image and needs OCR (a stretch goal).\n",
    "\n",
    "**What to notice:** real PDFs contain headers, page numbers and odd line breaks. That noise ends up in your chunks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2: Chunking, two ways\n",
    "**Why we do this:** a model cannot read a whole document per question, and one embedding of a whole document is too vague. We split text into *chunks*. Chunk strategy strongly affects retrieval quality (Module 3, section 5).\n",
    "\n",
    "**Key terms:**\n",
    "- *Fixed-size chunking:* cut every N characters, with an *overlap* so a sentence cut in half still appears whole in the next chunk.\n",
    "- *Sentence chunking:* group whole sentences up to a size limit, so no sentence is cut.\n",
    "\n",
    "**Your task:** complete `chunk_sentences`. Hint: keep adding sentences to the current chunk until adding another would exceed `size`, then save it and start a new one."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def chunk_fixed(pages, doc, size=800, overlap=150):\n",
    "    out = []\n",
    "    for page, text in pages:\n",
    "        text = re.sub(r\"\\s+\", \" \", text).strip()\n",
    "        step = size - overlap\n",
    "        for start in range(0, max(len(text), 1), step):\n",
    "            piece = text[start:start + size]\n",
    "            if piece.strip():\n",
    "                out.append({\"text\": piece, \"doc\": doc, \"page\": page})\n",
    "    return out\n",
    "\n",
    "def chunk_sentences(pages, doc, size=800):\n",
    "    out = []\n",
    "    for page, text in pages:\n",
    "        text = re.sub(r\"\\s+\", \" \", text).strip()\n",
    "        sentences = re.split(r\"(?<=[.!?])\\s+\", text)\n",
    "        current = \"\"\n",
    "        for s in sentences:\n",
    "            # YOUR CODE HERE: if current + s fits in `size`, append it to current;\n",
    "            # otherwise save `current` to `out` and start a new chunk with s.\n",
    "            pass\n",
    "        # YOUR CODE HERE: save the last `current` chunk if it is not empty\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details><summary>Solution (try first)</summary>\n",
    "\n",
    "```python\n",
    "for s in sentences:\n",
    "    if len(current) + len(s) + 1 <= size:\n",
    "        current = (current + \" \" + s).strip()\n",
    "    else:\n",
    "        if current:\n",
    "            out.append({\"text\": current, \"doc\": doc, \"page\": page})\n",
    "        current = s\n",
    "if current:\n",
    "    out.append({\"text\": current, \"doc\": doc, \"page\": page})\n",
    "```\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fixed = [c for d, p in DOCS.items() for c in chunk_fixed(p, d)]\n",
    "sent  = [c for d, p in DOCS.items() for c in chunk_sentences(p, d)]\n",
    "print(\"fixed chunks:\", len(fixed), \"| sentence chunks:\", len(sent))\n",
    "print(\"\\nSample fixed chunk:\\n\", fixed[3][\"text\"][:300])\n",
    "print(\"\\nSample sentence chunk:\\n\", sent[3][\"text\"][:300] if len(sent) > 3 else \"(complete the exercise)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**What to notice:** fixed chunks often begin or end mid-sentence. Sentence chunks read cleanly but vary in length."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3: Embed and store in ChromaDB\n",
    "**Why we do this:** an *embedding* turns text into a vector so that similar meanings sit close together. Chroma stores the vectors and finds the nearest ones to a question.\n",
    "\n",
    "**What the code does:** creates two collections, one per chunking strategy, so we can compare them. Chroma's built-in local embedding model is used, so no API key is needed (it downloads a small model on first use).\n",
    "\n",
    "**Finance link:** the data stays in memory on your machine. For real client documents this local option matters for confidentiality."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import chromadb\n",
    "client = chromadb.Client()\n",
    "\n",
    "def build(name, chunks):\n",
    "    col = client.get_or_create_collection(name)\n",
    "    col.add(\n",
    "        ids=[f\"{name}-{i}\" for i in range(len(chunks))],\n",
    "        documents=[c[\"text\"] for c in chunks],\n",
    "        metadatas=[{\"doc\": c[\"doc\"], \"page\": c[\"page\"]} for c in chunks],\n",
    "    )\n",
    "    return col\n",
    "\n",
    "col_fixed = build(\"fixed\", fixed)\n",
    "col_sent  = build(\"sentence\", sent) if sent else None\n",
    "print(\"stored:\", col_fixed.count(), \"fixed chunks\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 4: Retrieve and answer with citations\n",
    "**Why we do this:** this is the \"R\" and \"G\" of RAG. We retrieve the top-k chunks, put them in the prompt, and instruct the model to answer only from them and cite the source.\n",
    "\n",
    "**Key terms:** *top-k* = how many chunks to retrieve. *Grounding* = forcing the answer to rely on supplied text.\n",
    "\n",
    "**What the code does:** `retrieve` queries Chroma. `answer` builds a prompt with numbered sources and tells the model to say \"not found\" if the sources do not contain the answer.\n",
    "\n",
    "**Finance link:** the \"not found\" rule and citations are what make an answer auditable (Module 2, section 7)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def retrieve(col, question, k=4):\n",
    "    r = col.query(query_texts=[question], n_results=k)\n",
    "    return [{\"text\": t, **m} for t, m in zip(r[\"documents\"][0], r[\"metadatas\"][0])]\n",
    "\n",
    "SYSTEM = (\n",
    "    \"You answer questions about financial regulation using ONLY the numbered sources. \"\n",
    "    \"Cite sources like [1]. If the sources do not contain the answer, reply exactly: \"\n",
    "    \"'Not found in the provided documents.'\"\n",
    ")\n",
    "\n",
    "def answer(col, question, k=4):\n",
    "    hits = retrieve(col, question, k)\n",
    "    sources = \"\\n\\n\".join(f\"[{i+1}] ({h['doc']}, p.{h['page']})\\n{h['text']}\" for i, h in enumerate(hits))\n",
    "    reply = ask_llm(SYSTEM, f\"Sources:\\n{sources}\\n\\nQuestion: {question}\")\n",
    "    return reply, hits\n",
    "\n",
    "q = \"Who is responsible for the output of AI tools used by SEBI-regulated entities?\"\n",
    "reply, hits = answer(col_fixed, q)\n",
    "print(reply)\n",
    "print(\"\\nRetrieved from:\", [(h['doc'], h['page']) for h in hits])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** an answer citing `[1]`, `[2]` etc., stating that regulated persons stay responsible. Wording varies by model. Check each citation against the printed page numbers.\n",
    "\n",
    "**What to notice:** ask something the documents do *not* cover (e.g. \"What is the repo rate?\"). A well-grounded system says \"Not found\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 5: Compare chunking strategies\n",
    "**Why we do this:** you will measure, not guess, which chunking retrieves better.\n",
    "\n",
    "**What the code does:** for each test question you list a *keyword that must appear* in a good chunk. `hit_rate` is the share of questions where at least one of the top-k chunks contains the keyword. This is a crude proxy for *recall@k* (Module 3, section 6).\n",
    "\n",
    "**Your task:** write 5 questions about your documents with a distinctive keyword for each (fill in `TESTS`)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "TESTS = [\n",
    "    # (question, keyword that a correct chunk must contain)\n",
    "    (\"Who is solely responsible for AI tool output?\", \"responsib\"),\n",
    "    # YOUR CODE HERE: add 4 more (question, keyword) pairs from your documents\n",
    "]\n",
    "\n",
    "def hit_rate(col, k=4):\n",
    "    hits = 0\n",
    "    for q, kw in TESTS:\n",
    "        if any(kw.lower() in h[\"text\"].lower() for h in retrieve(col, q, k)):\n",
    "            hits += 1\n",
    "    return hits / len(TESTS)\n",
    "\n",
    "print(\"fixed   :\", hit_rate(col_fixed))\n",
    "if col_sent: print(\"sentence:\", hit_rate(col_sent))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**What to notice:** with only a handful of questions the numbers are noisy. A difference of one question is not conclusive. Try changing `size`, `overlap` and `k` and re-run.\n",
    "\n",
    "## Stretch goals\n",
    "1. Add the RBI and HDFC documents and test cross-document questions.\n",
    "2. Add metadata filtering (`where={\"doc\": ...}`) to restrict a query to one document.\n",
    "3. OCR a scanned PDF and add it to the corpus.\n",
    "\n",
    "## Common errors\n",
    "- **`ModuleNotFoundError`:** re-run the install cell, then restart the runtime.\n",
    "- **Connection refused (Ollama):** run `ollama serve` and `ollama pull llama3.1:8b`; use a local notebook, not Colab.\n",
    "- **Rate limit (Groq free tier):** wait a minute or lower `k`; limits change, check Groq's docs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}