{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Lab 4 Answer Key*\n",
    "\n",
    "# RAG Pipeline Design\n",
    "\n",
    "**Objective:** design a multi-tenant contract-RAG pipeline that serves both exact and semantic queries, enforces tenant isolation at retrieval time, and forces attribution.\n",
    "\n",
    "This is a design lab. The scenario: a legal-operations team wants Claude to answer questions over a multi-tenant corpus of contracts (PDFs containing prose, tables, and defined terms). Each contract belongs to exactly one tenant, and queries mix exact lookups (\"what does clause 7.2 say?\") with semantic ones (\"which of our contracts cap liability?\"). Most parts below are worked design artifacts: decision tables and rationales. Code appears only where the lab itself sketches code: structural chunking, retrieval fusion, and the groundedness check.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Confirm RAG is right\n",
    "\n",
    "The most expensive RAG mistake is building one you did not need. Run the subtraction test before you index anything: an exact key means a database lookup, a small stable corpus means a cached prefix, a computed answer means a tool. RAG is only for corpora too big to carry and queries too fuzzy to key.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Artifact: the subtraction test, applied to the scenario.**\n",
    "\n",
    "| Test | Applied to this corpus | Verdict |\n",
    "| --- | --- | --- |\n",
    "| Is the query an exact key? | Half the queries are exact (\"clause 7.2\"), but the other half are semantic (\"which contracts cap liability?\") | Not key-only: a database alone cannot serve the semantic half |\n",
    "| Small, stable, fits a cached prefix? | The corpus is large, growing, and multi-tenant; there is no single stable prefix any user could share | No: too big to carry, and tenancy forbids a shared prefix |\n",
    "| Is the answer computed, not retrieved? | Answers are readings of contract text, not totals or current state from a system of record | No: this is retrieval, not computation |\n",
    "| Too big to carry AND too fuzzy to key? | Yes on both halves of the definition | RAG qualifies |\n",
    "\n",
    "**The defensible paragraph.** RAG qualifies because the corpus is too big to carry and half the queries are too fuzzy to key. But the other half (\"what does clause 7.2 say?\") is an exact key, and dense embeddings famously miss rare tokens like clause numbers. So the pipeline must serve both query shapes rather than force everything through dense search. That single observation drives the retrieval design in Part 4: hybrid, not pure dense.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Parse and chunk structurally\n",
    "\n",
    "Chunk boundaries destroy meaning. A payment schedule split across two chunks answers no question, and a defined term split from its definition is worse than useless. Contracts have real structure (clauses), so the chunker must respect it, and metadata is born here: a chunk without metadata cannot be filtered, cited, or secured downstream.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Structural chunking: split on the document's own structure (clauses),\n",
    "# never on a 512-token grid, and attach the metadata every later\n",
    "# stage depends on. embed() is provider-agnostic pseudocode: Anthropic\n",
    "# ships no embedding endpoint; use your chosen third-party model.\n",
    "\n",
    "def chunk_contract(doc):\n",
    "    chunks = []\n",
    "    for clause in doc.iter_clauses():          # respect the real boundary\n",
    "        # Keep whole tables (header AND rows) inside their clause.\n",
    "        # A split liability schedule answers nothing.\n",
    "        text = clause.render(keep_tables_whole=True)\n",
    "        chunks.append({\n",
    "            \"text\": text,\n",
    "            \"vector\": embed(text),             # third-party embedding model\n",
    "            \"metadata\": {\n",
    "                \"source\": doc.id,\n",
    "                \"section\": clause.path,        # e.g. \"7.2 Limitation of Liability\"\n",
    "                \"version\": doc.version,        # staleness / re-index checks\n",
    "                \"tenant_id\": doc.tenant_id,    # ACL: enforced at retrieval (Part 5)\n",
    "                \"effective_date\": doc.effective_date,\n",
    "            },\n",
    "        })\n",
    "    return chunks\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Why each decision holds.** Fixed-size chunking on this corpus is an act of vandalism no reranker can undo: it splits payment schedules mid-table and severs defined terms from their definitions. Clause-level chunks are the document's own retrievable units. The refinement worth naming in a design review is parent-document retrieval: embed and retrieve the small precise unit (a clause) for accuracy, but feed the model the whole enclosing section at generation time so it has enough context to answer.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Index with metadata as a first-class citizen\n",
    "\n",
    "An index is not a bag of vectors. It is vectors plus the metadata that makes them governable: source, section, version, effective date, and tenant tags. The tenant_id and section fields are not optional; they are what make the next two parts possible.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Artifact: what each metadata field buys downstream.**\n",
    "\n",
    "| Field | What it enables | What breaks without it |\n",
    "| --- | --- | --- |\n",
    "| `source` | Citation and audit: every answer can name its document | Answers are unverifiable; Part 6 has nothing to check against |\n",
    "| `section` (e.g. \"7.2\") | Exact clause lookups via metadata filter; parent-document expansion | \"What does clause 7.2 say?\" falls through to dense search, which misses it |\n",
    "| `version` / `effective_date` | Recency filtering (filter first, then rank) and staleness detection after a document refresh | The system confidently answers from a superseded contract |\n",
    "| `tenant_id` | The ACL prefilter in Part 5 | Multi-tenancy is unenforceable; every query is a potential cross-tenant leak |\n",
    "\n",
    "One diagnostic rule rides along with `version`: if answers go wrong right after a document refresh while latency and model version hold steady, investigate the retrieval and indexing stage first (the data plane changed, nothing else did). And if the embedding model ever changes, re-embed the entire index: mixing two embedding spaces in one store silently returns garbage.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Hybrid retrieval matched to query shape\n",
    "\n",
    "Route or run both: lexical/BM25 catches \"clause 7.2\" and defined terms exactly; dense catches \"cap liability\" across different wording. Fuse the two ranked lists with reciprocal rank fusion, then rerank the survivors. Pure dense retrieval on identifiers is the classic wrong answer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hybrid retrieval: lexical (BM25) + dense, fused with reciprocal\n",
    "# rank fusion, ACL-filtered at retrieval time, then reranked.\n",
    "\n",
    "def reciprocal_rank_fusion(ranked_lists, k=60):\n",
    "    scores = {}\n",
    "    for ranked in ranked_lists:\n",
    "        for rank, chunk_id in enumerate(ranked):\n",
    "            scores[chunk_id] = scores.get(chunk_id, 0) + 1 / (k + rank)\n",
    "    return sorted(scores, key=scores.get, reverse=True)\n",
    "\n",
    "def retrieve(query, tenant_id, top_k=8):\n",
    "    lexical = bm25_search(query, limit=50)        # catches \"clause 7.2\", defined terms\n",
    "    dense = vector_search(embed(query), limit=50) # catches \"cap liability\", paraphrase\n",
    "    fused = reciprocal_rank_fusion([lexical, dense])\n",
    "\n",
    "    # ACL PREFILTER: security lives here, not in the prompt (Part 5).\n",
    "    allowed = [cid for cid in fused\n",
    "               if index.metadata(cid)[\"tenant_id\"] == tenant_id]\n",
    "\n",
    "    # Rerank the ACL-safe survivors; keep only the true top-k.\n",
    "    # A tight top-k beats a generous retriever: monolithic context\n",
    "    # degrades when the window fills with marginal chunks.\n",
    "    return rerank(query, allowed)[:top_k]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Design notes to defend out loud.** The fusion step is what lets one pipeline serve two query shapes without a brittle router. The reranker resolves the merged list so the best chunk does not sit at rank 20. Keep top-k tight: more retrieved chunks is not more accuracy, and padding the window with marginally relevant clauses lowers answer quality while raising cost. Pair this with parent-document expansion so a clause hit carries its whole section into context.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Enforce the ACL at retrieval, not in the prompt\n",
    "\n",
    "Apply tenant_id as a prefilter before rerank and before context assembly. A rival tenant's clause must never become a candidate. Do not, under any circumstances, retrieve across tenants and ask Claude to ignore the wrong ones.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The rationale, in the form a security review expects.**\n",
    "\n",
    "The naive design retrieves top-k across all tenants and appends \"only use documents belonging to the requesting user's organization.\" That is not a security control; it is a suggestion to a text generator, and it fails the moment the model is confused or the prompt is injected. The document that must never reach the user must never reach the model.\n",
    "\n",
    "**Where the control lives in this design:**\n",
    "\n",
    "1. `tenant_id` is attached at chunking time (Part 2), so every chunk carries its owner.\n",
    "2. The prefilter in `retrieve()` (Part 4) drops cross-tenant chunks after fusion and before rerank, so a rival tenant's clause is never a rerank candidate, never assembled into context, and never one generation slip away from the response.\n",
    "3. The tenant identity comes from the authenticated request, never from model output or user text, so an injected instruction cannot widen the filter.\n",
    "\n",
    "**The one-sentence version for the design review:** multi-tenancy is enforced by the index and the query, not by the prompt.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6: Force attribution and check groundedness\n",
    "\n",
    "An answer you cannot attribute is an answer you cannot verify. Constrain generation to return the answer and the IDs of the chunks it cited, then confirm every cited chunk was actually in the retrieved set. This produces the audit trail the legal team needs and the labelled data the eval loop (Lab 6) consumes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import anthropic\n",
    "\n",
    "client = anthropic.Anthropic()\n",
    "\n",
    "ANSWER_SCHEMA = {\n",
    "    \"type\": \"json_schema\",\n",
    "    \"schema\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"answer\": {\"type\": \"string\"},\n",
    "            # The model must name the chunks it used. Downstream code\n",
    "            # rejects any id that was not in the retrieved set.\n",
    "            \"cited_chunk_ids\": {\n",
    "                \"type\": \"array\",\n",
    "                \"items\": {\"type\": \"string\"},\n",
    "            },\n",
    "        },\n",
    "        \"required\": [\"answer\", \"cited_chunk_ids\"],\n",
    "        \"additionalProperties\": False,\n",
    "    },\n",
    "}\n",
    "\n",
    "response = client.messages.create(\n",
    "    model=\"claude-opus-4-8\",\n",
    "    max_tokens=4096,\n",
    "    thinking={\"type\": \"adaptive\"},\n",
    "    output_config={\"effort\": \"high\", \"format\": ANSWER_SCHEMA},\n",
    "    system=[{\n",
    "        \"type\": \"text\",\n",
    "        \"text\": \"Answer only from the provided context. Every claim must \"\n",
    "                \"cite the chunk id it came from. If the context does not \"\n",
    "                \"contain the answer, say so.\",\n",
    "        \"cache_control\": {\"type\": \"ephemeral\"},\n",
    "    }],\n",
    "    messages=[{\"role\": \"user\", \"content\": build_context(retrieved_chunks, question)}],\n",
    ")\n",
    "\n",
    "# Groundedness check: every cited id must exist in what we retrieved.\n",
    "class GroundednessError(Exception):\n",
    "    pass\n",
    "\n",
    "result = parse(response)\n",
    "retrieved_ids = {c[\"id\"] for c in retrieved_chunks}\n",
    "if not set(result[\"cited_chunk_ids\"]).issubset(retrieved_ids):\n",
    "    raise GroundednessError(\"model cited a chunk it was never given\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Why this closes the design.** The schema is forced through `output_config.format` (with `effort` nested alongside it), and `thinking` is adaptive. The groundedness check is what turns a citation from decoration into a control: a cited id outside the retrieved set is rejected before anything reaches the user. With attribution in place, every answer becomes a labelled, auditable event: the eval loop (Lab 6) gets its data, and the audit trail (Lab 8) gets its evidence. The exercise compresses the whole lab: subtraction test, structural chunking, metadata-rich index, hybrid retrieval, ACL at retrieval, forced attribution.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
