L2 · 06 · RAG: Retrieval-Augmented Generation

Models don't know your company's return policy, store procedures or internal docs, and they will happily invent them. RAG fixes this: retrieve the relevant passages, put them in the prompt, and make the model answer only from them, with citations. This page builds a policy assistant for store staff.

Why RAG (and not fine-tuning) for knowledge

The two pipelines

Chunking

You retrieve chunks, not whole documents. Chunk quality matters more than almost any other RAG decision.

• Prefix each chunk with its context, such as "Returns Policy › Opened items:". A chunk saying "within 14 days" is useless without knowing what it refers to.
• Store metadata: source document, section, URL, last-updated date, audience/permissions.
• Too small loses context. Too big dilutes relevance and wastes tokens. Test a few sizes.

The grounded prompt

You answer questions from retail store staff using ONLY the sources below.
Rules:
- Cite sources inline like [1] or [2][3] after each claim.
- If the sources do not contain the answer, reply exactly: "I can't find this in the policies."
- Do not use outside knowledge. Do not guess numbers, dates or amounts.
- Be concise: at most 4 sentences.

<sources>
[1] (Returns Policy > Standard returns) Most items can be returned within 14 days of purchase or pickup with a receipt or order ID.
[2] (Returns Policy > Opened items) Opened headphones and earbuds can be returned within 14 days only if defective.
[3] (Pickup > Holding period) Pickup orders are held for 7 days, then cancelled and refunded.
</sources>

Question: A customer opened earbuds 10 days ago and doesn't like the sound. Can they return them?

Expected answer: "Not for that reason. Opened earbuds can only be returned within 14 days if they are defective [2]. Disliking the sound isn't a defect." The citation lets staff check the source, and lets your code verify that the model cited a real source.

Scala 3: a minimal policy RAG

Runs offline with the toy embedder from page 05 and a scripted generator that prints the exact prompt that would be sent. Swap in the live embedder and the LLM client from page 03 for real use.

//> using scala 3

final case class Chunk(id: Int, source: String, section: String, text: String):
  def forEmbedding: String = s"$source > $section: $text"

// ---------- Ingestion ----------
val docs = Map(
  "Returns Policy" ->
    """## Standard returns
      |Most items can be returned within 14 days of purchase or pickup with a receipt or order ID.
      |Refunds go to the original payment method within 5 business days.
      |## Opened items
      |Opened headphones and earbuds can be returned within 14 days only if defective.
      |Opened software and gift cards cannot be returned.""".stripMargin,
  "Pickup" ->
    """## Holding period
      |Pickup orders are held for 7 days, then cancelled and refunded.
      |## Identification
      |Customers must show the order confirmation and a photo ID matching the order name.""".stripMargin
)

/** Structure-aware chunking: one chunk per "## section". */
def chunk(docs: Map[String, String]): Vector[Chunk] =
  docs.toVector.sortBy(_._1).flatMap { (source, body) =>
    body.split("(?m)^## ").toVector.filter(_.trim.nonEmpty).map { sec =>
      val lines = sec.linesIterator.toVector
      (source, lines.head.trim, lines.tail.mkString(" ").trim)
    }
  }.zipWithIndex.map { case ((src, sec, txt), i) => Chunk(i + 1, src, sec, txt) }

// ---------- Embedding (toy, offline; see page 05 for the live version) ----------
type Vec = Array[Double]
def embed(t: String): Vec =
  val v = new Array[Double](512)
  val words = t.toLowerCase.split("\\W+").filter(_.length > 2)
  words.foreach(w => v(math.floorMod(w.hashCode, 512)) += 1)
  val n = math.sqrt(v.map(x => x * x).sum)
  if n == 0 then v else v.map(_ / n)
def cosine(a: Vec, b: Vec): Double = a.indices.map(i => a(i) * b(i)).sum

final class Retriever(chunks: Vector[Chunk]):
  private val vectors = chunks.map(c => embed(c.forEmbedding))
  def topK(question: String, k: Int = 3, minScore: Double = 0.05): Vector[(Chunk, Double)] =
    val q = embed(question)
    chunks.zip(vectors).map((c, v) => c -> cosine(q, v)).filter(_._2 >= minScore).sortBy(-_._2).take(k)

// ---------- Augment ----------
def buildPrompt(question: String, hits: Vector[Chunk]): String =
  val sources = hits.zipWithIndex.map((c, i) => s"[${i + 1}] (${c.source} > ${c.section}) ${c.text}").mkString("\n")
  s"""You answer questions from retail store staff using ONLY the sources below.
     |Cite sources inline like [1]. If the sources do not contain the answer, reply exactly:
     |"I can't find this in the policies." Do not guess. At most 4 sentences.
     |
     |<sources>
     |$sources
     |</sources>
     |
     |Question: $question""".stripMargin

// ---------- Check citations ----------
def citationProblems(answer: String, numSources: Int): List[String] =
  val cited = """\[(\d+)\]""".r.findAllMatchIn(answer).map(_.group(1).toInt).toSet
  val invalid = cited.filter(n => n < 1 || n > numSources).toList.sorted.map(n => s"cites non-existent source [$n]")
  val missing = Option.when(cited.isEmpty && !answer.contains("can't find"))("answer has no citations").toList
  invalid ++ missing

@main def policyRag(): Unit =
  val retriever = Retriever(chunk(docs))
  val question  = "A customer opened earbuds 10 days ago and does not like the sound. Can they return them?"
  val hits      = retriever.topK(question)
  hits.foreach((c, s) => println(f"retrieved $s%.2f  #${c.id} ${c.source} > ${c.section}"))
  println("\n----- prompt sent to the model -----\n" + buildPrompt(question, hits.map(_._1)))

  // Scripted answers to demonstrate the citation check (use a real model in practice)
  for answer <- List(
      "Not for that reason. Opened earbuds can only be returned within 14 days if defective [1].",
      "Yes, all items can be returned within 30 days [4]."
    ) do
    val problems = citationProblems(answer, hits.size)
    println(s"\nanswer: $answer\ncheck : ${if problems.isEmpty then "OK" else problems.mkString("; ")}")

Python: the same pipeline

import math
import re

DOCS = {
    "Returns Policy": "## Standard returns\nMost items can be returned within 14 days of purchase or pickup with a receipt or order ID.\n"
                      "## Opened items\nOpened headphones and earbuds can be returned within 14 days only if defective.",
    "Pickup": "## Holding period\nPickup orders are held for 7 days, then cancelled and refunded.\n"
              "## Identification\nCustomers must show the order confirmation and a photo ID matching the order name.",
}


def chunk(docs: dict[str, str]) -> list[dict]:
    chunks = []
    for source, body in sorted(docs.items()):
        for sec in filter(str.strip, re.split(r"(?m)^## ", body)):
            title, _, text = sec.partition("\n")
            chunks.append({"source": source, "section": title.strip(), "text": " ".join(text.split())})
    return chunks


def embed(text: str, dims: int = 512) -> list[float]:   # toy; use a real embedding model in practice
    v = [0.0] * dims
    for w in re.findall(r"\w{3,}", text.lower()):
        v[hash(w) % dims] += 1
    n = math.sqrt(sum(x * x for x in v)) or 1.0
    return [x / n for x in v]


CHUNKS = chunk(DOCS)
VECTORS = [embed(f"{c['source']} > {c['section']}: {c['text']}") for c in CHUNKS]


def retrieve(question: str, k: int = 3) -> list[dict]:
    q = embed(question)
    scored = sorted(zip(CHUNKS, VECTORS), key=lambda cv: -sum(a * b for a, b in zip(q, cv[1])))
    return [c for c, _ in scored[:k]]


def build_prompt(question: str, hits: list[dict]) -> str:
    sources = "\n".join(f"[{i}] ({c['source']} > {c['section']}) {c['text']}" for i, c in enumerate(hits, 1))
    return ("Answer using ONLY the sources. Cite like [1]. If not covered, reply exactly: "
            f"\"I can't find this in the policies.\"\n\n<sources>\n{sources}\n</sources>\n\nQuestion: {question}")


def citation_problems(answer: str, n_sources: int) -> list[str]:
    cited = {int(n) for n in re.findall(r"\[(\d+)\]", answer)}
    problems = [f"cites non-existent source [{n}]" for n in sorted(cited) if not 1 <= n <= n_sources]
    if not cited and "can't find" not in answer:
        problems.append("answer has no citations")
    return problems


if __name__ == "__main__":
    q = "How long do we keep a pickup order before cancelling it?"
    hits = retrieve(q)
    print(build_prompt(q, hits))
    print(citation_problems("Pickup orders are held for 7 days [1].", len(hits)))   # []
    print(citation_problems("They are held for 30 days.", len(hits)))              # ['answer has no citations']

Evaluating a RAG system

Evaluate the two halves separately, because they fail for different reasons.

Start with 30–50 real staff questions, each with the expected source chunk and a reference answer. Re-run after every change to chunking, embeddings, k or the prompt.

Common failure modes