{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Chapter 7 ยท Build a RAG Pipeline ๐Ÿ”Ž\n", "\n", "**Retrieval-Augmented Generation, hands-on.**\n", "\n", "An LLM only knows what it was trained on โ€” not last week's news, and *definitely* not your company's private documents. You can't cram 500 GB into a prompt. So instead of baking knowledge into the model, you **retrieve it the moment you need it**. That's RAG.\n", "\n", "Think of it like an open-book exam. The model isn't trying to remember everything from memory โ€” you hand it the *right page* of the textbook right before it answers.\n", "\n", "RAG is **three phases**:\n", "\n", "| Phase | When | What happens |\n", "|-------|------|--------------|\n", "| ๐Ÿ—„๏ธ **INDEX** | once, ahead of time | split docs into chunks โ†’ turn each into a vector โ†’ store them |\n", "| ๐Ÿงฒ **RETRIEVE** | per question | embed the question โ†’ find the nearest chunk vectors |\n", "| ๐Ÿค– **GENERATE** | per question | drop those chunks into the prompt โ†’ answer, grounded |\n", "\n", "In this notebook you'll build all three against a tiny fictional company's policy docs, then watch the payoff: a **grounded** answer vs. the model **guessing** without RAG.\n", "\n", "> Everything here runs on Colab's free tier: a local **Chroma** vector store + free **Gemini** embeddings + Gemini for the answer." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Setup ๐Ÿ› ๏ธ\n", "\n", "Two libraries:\n", "- **`google-genai`** โ€” Google's official SDK. We use it for *both* embeddings and text generation.\n", "- **`chromadb`** โ€” a lightweight vector database that runs entirely in memory, right here in the notebook. No server, no signup.\n", "\n", "Run the cell below (takes ~30 seconds)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install -q google-genai chromadb" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Your API key ๐Ÿ”‘\n", "\n", "You need a free Gemini API key. Grab one here ๐Ÿ‘‰ **https://aistudio.google.com/apikey** (takes 20 seconds, no credit card).\n", "\n", "**In Colab**, the clean way is the ๐Ÿ”‘ **Secrets** panel in the left sidebar:\n", "1. Click the key icon, add a secret named `GOOGLE_API_KEY`, paste your key.\n", "2. Toggle *Notebook access* on.\n", "\n", "The cell below reads it from there, and falls back to a hidden prompt if you haven't set it up." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "KEY = None\n", "try:\n", " from google.colab import userdata\n", " KEY = userdata.get(\"GOOGLE_API_KEY\")\n", "except Exception:\n", " pass\n", "\n", "if not KEY:\n", " KEY = os.environ.get(\"GOOGLE_API_KEY\")\n", "\n", "if not KEY:\n", " from getpass import getpass\n", " KEY = getpass(\"Paste your Gemini API key (get one at https://aistudio.google.com/apikey): \")\n", "\n", "assert KEY, \"No API key found โ€” set GOOGLE_API_KEY or paste it above.\"\n", "print(\"Key loaded โœ…\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### One client, two jobs\n", "\n", "We create a single `genai.Client`. We'll ask it to do two different things:\n", "- `client.models.embed_content(...)` โ†’ turn text into vectors (the *fingerprint of meaning*)\n", "- `client.models.generate_content(...)` โ†’ write the final answer\n", "\n", "โš ๏ธ **The one rule that matters most in RAG:** you must embed your documents *and* your questions with the **same embedding model**. Vectors from different models live in different \"coordinate systems\" โ€” comparing them is like comparing latitude to Fahrenheit. We pin ours to `EMB_MODEL` and reuse it everywhere." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from google import genai\n", "\n", "client = genai.Client(api_key=KEY)\n", "\n", "EMB_MODEL = \"text-embedding-004\" # SAME model for indexing AND querying\n", "GEN_MODEL = \"gemini-2.0-flash\" # writes the grounded answer\n", "\n", "\n", "def embed(texts):\n", " \"\"\"texts: list[str] -> list[list[float]] (one vector per text).\"\"\"\n", " try:\n", " # Preferred: batch โ€” send all texts in one call.\n", " r = client.models.embed_content(model=EMB_MODEL, contents=texts)\n", " return [e.values for e in r.embeddings]\n", " except Exception:\n", " # Defensive fallback: some SDK versions want one item at a time.\n", " out = []\n", " for t in texts:\n", " r = client.models.embed_content(model=EMB_MODEL, contents=t)\n", " out.append(r.embeddings[0].values)\n", " return out\n", "\n", "\n", "# quick smoke test\n", "v = embed([\"hello world\"])\n", "print(f\"Got {len(v)} vector, {len(v[0])} dimensions. First few numbers: {v[0][:5]}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Sample documents ๐Ÿ“„\n", "\n", "Meet **Northwind Softworks**, a fictional SaaS company. Below are six short snippets from their support wiki โ€” refund rules, shipping, support hours. This is our \"private knowledge\" that the model was never trained on.\n", "\n", "Notice there's a *deliberate* trap: refund rules differ for **annual** vs **monthly** plans. Good retrieval has to pull the *right* one." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "documents = [\n", " \"Refund policy for annual plans: Annual subscriptions are fully refundable within 30 days \"\n", " \"of purchase for a 100% refund. After 30 days, refunds are prorated based on the unused \"\n", " \"months remaining in the term. To request a refund, contact billing@northwind.example.\",\n", "\n", " \"Refund policy for monthly plans: Monthly subscriptions are non-refundable once the billing \"\n", " \"cycle has started, but you can cancel anytime to stop all future charges. Cancellation takes \"\n", " \"effect at the end of the current paid month.\",\n", "\n", " \"Shipping: Physical welcome kits and branded hardware ship free within the continental US and \"\n", " \"arrive in 5-7 business days. International shipping costs $18 flat and takes 2-3 weeks. \"\n", " \"We do not ship to P.O. boxes.\",\n", "\n", " \"Support hours: Live chat and phone support run Monday through Friday, 8am to 6pm Central Time. \"\n", " \"Weekend coverage is email-only, with responses within 24 hours. Enterprise customers get \"\n", " \"24/7 priority support.\",\n", "\n", " \"Data & security: All customer data is encrypted at rest and in transit. Northwind is SOC 2 \"\n", " \"Type II certified and stores data in US-based data centers. Customers can request a full \"\n", " \"data export or permanent deletion at any time.\",\n", "\n", " \"Free trial: New accounts get a 14-day free trial of the Pro plan with no credit card required. \"\n", " \"When the trial ends, accounts automatically downgrade to the Free tier unless you choose a \"\n", " \"paid plan. No charges happen without explicit confirmation.\",\n", "]\n", "\n", "print(f\"{len(documents)} documents loaded.\")\n", "for i, d in enumerate(documents):\n", " print(f\"[{i}] {d[:70]}...\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. INDEX ๐Ÿ—„๏ธ โ€” chunk โ†’ embed โ†’ store\n", "\n", "This is the \"once, ahead of time\" phase. Three moves:\n", "\n", "1. **Chunk** โ€” split documents into small pieces. Our docs are already one-idea-each paragraphs, so here each document *is* a chunk. (In the real world you'd slice long PDFs into ~300-500 character windows โ€” more on that below.)\n", "2. **Embed** โ€” turn every chunk into a vector using our pinned `EMB_MODEL`.\n", "3. **Store** โ€” drop the chunks + their vectors into Chroma so we can search them later.\n", "\n", "The beauty: you pay this cost *once*. Query a thousand times afterward and you never re-embed the documents." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# 1. CHUNK โ€” one paragraph = one chunk (simple + clean for this dataset)\n", "chunks = documents\n", "ids = [f\"chunk-{i}\" for i in range(len(chunks))]\n", "\n", "# 2. EMBED โ€” every chunk becomes a vector (with the SAME model we'll query with)\n", "chunk_vectors = embed(chunks)\n", "print(f\"Embedded {len(chunk_vectors)} chunks into {len(chunk_vectors[0])}-dim vectors.\")\n", "\n", "# 3. STORE โ€” put them in a fresh Chroma collection\n", "import chromadb\n", "\n", "chroma = chromadb.Client()\n", "\n", "# start clean so re-running this cell doesn't error on a duplicate name\n", "try:\n", " chroma.delete_collection(\"policies\")\n", "except Exception:\n", " pass\n", "\n", "col = chroma.create_collection(\"policies\")\n", "col.add(ids=ids, documents=chunks, embeddings=chunk_vectors)\n", "\n", "print(f\"Stored {col.count()} chunks in the 'policies' collection โœ…\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. RETRIEVE ๐Ÿงฒ โ€” embed the question, find the nearest chunks\n", "\n", "Now the per-question phase begins. The key idea: **similarity search, not keyword matching.** We embed the *question* with the exact same model, then ask Chroma for the chunks whose vectors sit closest to it.\n", "\n", "Watch the `distances` in the output: **smaller = more similar**. The annual-refund chunk should win clearly, even though our question never uses the word \"subscription.\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "question = \"What's the refund policy for annual plans?\"\n", "\n", "# embed the question the SAME way we embedded the chunks\n", "q_vector = embed([question])\n", "\n", "# ask Chroma for the 3 nearest chunks\n", "hits = col.query(query_embeddings=q_vector, n_results=3)\n", "\n", "print(f\"Question: {question}\\n\")\n", "print(\"Retrieved chunks (nearest first):\\n\")\n", "for doc, dist in zip(hits[\"documents\"][0], hits[\"distances\"][0]):\n", " print(f\" distance={dist:.3f} | {doc[:90]}...\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "See how the **annual** refund chunk comes back first, and the **monthly** one is farther away? That gap is retrieval doing its job โ€” it understood *meaning*, not just matching letters. That top chunk is what we'll hand to the model next." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. GENERATE ๐Ÿค– โ€” a grounded answer\n", "\n", "Final phase. We stitch the retrieved chunks into a prompt with a strict instruction: **answer using ONLY this context.** Then we send it to Gemini.\n", "\n", "This is the whole trick. The model isn't recalling facts from training โ€” it's *reading* the pages we handed it and summarizing. Current, private, sourced." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# stitch the retrieved chunks into a context block\n", "retrieved = hits[\"documents\"][0]\n", "context = \"\\n\\n\".join(retrieved)\n", "\n", "prompt = f\"\"\"Answer the question using ONLY the context below.\n", "If the answer is not in the context, say you don't have that information.\n", "\n", "Context:\n", "{context}\n", "\n", "Question: {question}\n", "Answer:\"\"\"\n", "\n", "resp = client.models.generate_content(model=GEN_MODEL, contents=prompt)\n", "print(\"๐Ÿค– Grounded answer:\\n\")\n", "print(resp.text)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6. The payoff ๐ŸŽฏ โ€” grounded vs. guessing\n", "\n", "Here's why any of this matters. Let's ask the model the **exact same question** with **no context at all** โ€” just the raw question. Northwind is fictional; its policies were never in the training data. So the model has three bad options: refuse, hedge, or *make something up*.\n", "\n", "Run it and compare the two answers side by side." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# NO retrieval, no context โ€” just the bare question\n", "naked_prompt = question\n", "\n", "no_rag = client.models.generate_content(model=GEN_MODEL, contents=naked_prompt)\n", "\n", "print(\"=\" * 60)\n", "print(\"โŒ WITHOUT RAG (no context โ€” the model is on its own):\")\n", "print(\"=\" * 60)\n", "print(no_rag.text)\n", "print()\n", "print(\"=\" * 60)\n", "print(\"โœ… WITH RAG (grounded in Northwind's actual docs):\")\n", "print(\"=\" * 60)\n", "print(resp.text)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**That contrast is the entire value of RAG.**\n", "\n", "- Without context, the model has never heard of Northwind โ€” so it either says it can't help, or invents a plausible-sounding policy (a *hallucination*). Neither is safe to show a customer.\n", "- With context, the answer is specific, correct, and traceable to a source document.\n", "\n", "Same model. Same question. The only difference is whether we handed it the right page first." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### ๐Ÿ“ A note on chunking\n", "\n", "We cheated slightly: our documents were already bite-sized. Real documents aren't. **How you chunk decides how well RAG works** โ€” it's the make-or-break craft:\n", "\n", "- **Too small** โ†’ precise matches, but each chunk is stranded without its surrounding context. (\"Refunds are prorated.\" Prorated *for what plan*? The chunk doesn't say.)\n", "- **Too large** โ†’ full context, but fuzzy retrieval. The one relevant sentence hides inside paragraphs of noise, so similarity gets diluted.\n", "- **The sweet spot** is usually ~300-500 characters (or a few sentences) with a little overlap between neighbors so no idea gets cut in half.\n", "\n", "A clever pro move is **hierarchical chunking**: match on tiny precise chunks, but hand the model the *larger parent* chunk for context โ€” precision of small, completeness of large." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## ๐ŸŽฏ Your turn\n", "\n", "You've built the full pipeline. Now make it yours:\n", "\n", "1. **Ask harder questions.** Try `\"Can I get money back on a monthly plan?\"` โ€” does retrieval correctly grab the *monthly* chunk instead of the annual one? Try one whose answer *isn't* in the docs (e.g. `\"Do you offer student discounts?\"`) and confirm the grounded answer admits it doesn't know, while the no-RAG version guesses.\n", "\n", "2. **Tune `n_results`.** Change it to `1` and to `5`. Fewer chunks = tighter but riskier (you might miss the answer). More chunks = safer but noisier (and a longer prompt). Find the trade-off.\n", "\n", "3. **Bring your own docs.** Replace the `documents` list with real text โ€” a few paragraphs from a syllabus, a game manual, your own notes โ€” then re-run INDEX and ask questions about it. This is the exact moment RAG stops being a demo and becomes *useful*.\n", "\n", "4. **Stretch:** write a `rag(question)` function that wraps all three phases (retrieve โ†’ build prompt โ†’ generate) into a single call, so you can ask questions in a loop.\n", "\n", "When you can point RAG at your own documents and trust the answers, you've got the single most important practical AI skill there is. ๐Ÿš€" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "Part of the Agentic AI course ยท Chapter 7" ] } ], "metadata": { "colab": { "provenance": [], "name": "ch7-rag-pipeline.ipynb" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }