RAG starter with hybrid search and recall at k

1"""A small, complete RAG pipeline you can read in one sitting and run on a laptop.
2
3What it does, in order:
4  1. Chunks markdown documents by heading, keeping the document title and heading in every chunk.
5  2. Indexes the chunks twice: BM25 over words, and character n-gram vectors (a cheap, offline stand-in for
6     embeddings that catches word forms and typos).
7  3. Retrieves with each index and fuses the two rankings with reciprocal rank fusion (hybrid search).
8  4. Answers with the best supported sentences and a citation for each, or says it does not know when the
9     evidence is weak. With --llm it asks a model instead, given only the retrieved chunks.
10
11Standard library only, except the optional --llm mode, which needs the `openai` package and an
12OpenAI-compatible endpoint (OPENAI_API_KEY, optional OPENAI_BASE_URL for a local server, RAG_MODEL for the model).
13
14    python rag.py "How much can I spend on a hotel in London?"
15    python rag.py --mode bm25 "Can I work from Spain for a month?"
16"""
17
18from __future__ import annotations
19
20import argparse
21import math
22import os
23import re
24import sys
25from collections import Counter
26from dataclasses import dataclass
27from pathlib import Path
28
29CORPUS = Path(__file__).parent / "corpus"
30
31STOPWORDS = set(
32    "a an and are as at be but by can do does for from has have how i if in into is it its me my no not of on or "
33    "our so than that the their then there these they this to up us was we were what when where which who why will "
34    "with you your".split()
35)
36
37
38@dataclass
39class Chunk:
40    id: str  # "<document>#<heading>", the citation a reader can follow
41    title: str
42    heading: str
43    text: str
44
45    @property
46    def indexed(self) -> str:
47        # The heading path goes into the indexed text: "Hotels" alone says little; "Travel and expenses > Hotels" says where.
48        return f"{self.title} > {self.heading}\n{self.text}" if self.heading else self.text
49
50
51def slug(text: str) -> str:
52    return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
53
54
55def load_chunks(folder: Path = CORPUS) -> list[Chunk]:
56    """One chunk per "## " section. Splitting on the document's own structure keeps a rule and its exceptions together."""
57    chunks: list[Chunk] = []
58    for path in sorted(folder.glob("*.md")):
59        title, heading, lines = path.stem, "", []
60
61        def flush() -> None:
62            body = " ".join(line.strip() for line in lines if line.strip())
63            if heading and body:
64                chunks.append(Chunk(f"{path.stem}#{slug(heading)}", title, heading, body))
65
66        for line in path.read_text(encoding="utf-8").splitlines():
67            if line.startswith("# "):
68                title = line[2:].strip()
69            elif line.startswith("## "):
70                flush()
71                heading, lines = line[3:].strip(), []
72            else:
73                lines.append(line)
74        flush()
75    return chunks
76
77
78def load_windows(folder: Path = CORPUS, size: int = 40) -> list[Chunk]:
79    """The common alternative, for comparison: fixed windows of `size` words that ignore headings.
80
81    Each window is labelled with the section it starts in, so the evaluation can still score it. Text that runs
82    across a boundary lands in a window labelled with the previous section, which is exactly the problem.
83    """
84    chunks: list[Chunk] = []
85    for path in sorted(folder.glob("*.md")):
86        words: list[str] = []
87        owners: list[str] = []
88        heading = ""
89        for line in path.read_text(encoding="utf-8").splitlines():
90            if line.startswith("# "):
91                continue
92            if line.startswith("## "):
93                heading = line[3:].strip()
94                continue
95            for word in line.split():
96                words.append(word)
97                owners.append(heading)
98        for start in range(0, len(words), size):
99            # No title or heading goes into a window's indexed text: that context is what fixed windows lose.
100            chunks.append(Chunk(f"{path.stem}#{slug(owners[start])}", "", "", " ".join(words[start : start + size])))
101    return chunks
102
103
104def stem(word: str) -> str:
105    """Rough suffix stripping, enough for "holidays" to meet "holiday" and "claims" to meet "claim"."""
106    for suffix in ("ing", "ed", "es", "s"):
107        if len(word) > len(suffix) + 2 and word.endswith(suffix):
108            return word[: -len(suffix)]
109    return word
110
111
112def terms(text: str) -> list[str]:
113    return [stem(w) for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOPWORDS]
114
115
116class BM25:
117    """Okapi BM25: rewards chunks that contain the query's rare words, without letting long chunks win by length."""
118
119    def __init__(self, docs: list[str], k1: float = 1.5, b: float = 0.75) -> None:
120        self.k1, self.b = k1, b
121        self.tf = [Counter(terms(d)) for d in docs]
122        self.lengths = [sum(tf.values()) for tf in self.tf]
123        self.avg = sum(self.lengths) / len(self.lengths)
124        df = Counter(t for tf in self.tf for t in tf)
125        n = len(docs)
126        self.idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()}
127
128    def scores(self, query: str) -> list[float]:
129        q = terms(query)
130        out = []
131        for tf, length in zip(self.tf, self.lengths):
132            s = 0.0
133            for t in q:
134                f = tf.get(t, 0)
135                if f:
136                    s += self.idf[t] * f * (self.k1 + 1) / (f + self.k1 * (1 - self.b + self.b * length / self.avg))
137            out.append(s)
138        return out
139
140
141class CharNgrams:
142    """TF-IDF over character trigrams inside words, compared by cosine similarity.
143
144    It catches word forms and spelling slips that exact words miss ("reimburse" and "reimbursed", "postmortem" and
145    "post-mortem"), but it does not understand meaning. Swap in an embedding model here when you have one: the
146    rest of the pipeline only needs one score per chunk.
147    """
148
149    def __init__(self, docs: list[str], n: int = 3) -> None:
150        self.n = n
151        counts = [self._grams(d) for d in docs]
152        df = Counter(g for c in counts for g in c)
153        self.idf = {g: math.log(len(docs) / f) + 1 for g, f in df.items()}
154        self.vectors = [self._weigh(c) for c in counts]
155
156    def _grams(self, text: str) -> Counter:
157        grams: Counter = Counter()
158        for word in terms(text):
159            padded = f" {word} "
160            grams.update(padded[i : i + self.n] for i in range(len(padded) - self.n + 1))
161        return grams
162
163    def _weigh(self, counts: Counter) -> dict[str, float]:
164        vec = {g: c * self.idf.get(g, 0.0) for g, c in counts.items()}
165        norm = math.sqrt(sum(v * v for v in vec.values())) or 1.0
166        return {g: v / norm for g, v in vec.items()}
167
168    def scores(self, query: str) -> list[float]:
169        q = self._weigh(self._grams(query))
170        return [sum(w * vec.get(g, 0.0) for g, w in q.items()) for vec in self.vectors]
171
172
173def rank(scores: list[float]) -> list[int]:
174    return sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
175
176
177def reciprocal_rank_fusion(rankings: list[list[int]], k: int = 60) -> list[int]:
178    """Merge rankings by position, not by score: BM25 scores and cosines live on different scales."""
179    fused: Counter = Counter()
180    for ranking in rankings:
181        for position, i in enumerate(ranking):
182            fused[i] += 1 / (k + position + 1)
183    return [i for i, _ in fused.most_common()]
184
185
186class Retriever:
187    def __init__(self, chunks: list[Chunk]) -> None:
188        self.chunks = chunks
189        docs = [c.indexed for c in chunks]
190        self.bm25 = BM25(docs)
191        self.ngrams = CharNgrams(docs)
192
193    def search(self, query: str, k: int = 5, mode: str = "hybrid") -> list[Chunk]:
194        by_words = rank(self.bm25.scores(query))
195        by_grams = rank(self.ngrams.scores(query))
196        order = {"bm25": by_words, "ngram": by_grams, "hybrid": reciprocal_rank_fusion([by_words, by_grams])}[mode]
197        return [self.chunks[i] for i in order[:k]]
198
199
200def coverage(query: str, text: str) -> float:
201    """Share of the question's meaningful words that appear in a chunk: a crude, explainable evidence check."""
202    q = set(terms(query))
203    return len(q & set(terms(text))) / len(q) if q else 0.0
204
205
206def extractive_answer(query: str, hits: list[Chunk], threshold: float = 0.4) -> str:
207    """Quote the best supported sentences with their citation, or decline when the evidence is too thin."""
208    best = max(hits, key=lambda c: coverage(query, c.indexed), default=None)
209    if best is None or coverage(query, best.indexed) < threshold:
210        closest = ", ".join(c.id for c in hits[:3])
211        return f"I could not find this in the documents. Closest sections: {closest}."
212    sentences = re.split(r"(?<=[.!?])\s+", best.text)
213    picked = sorted(sentences, key=lambda s: coverage(query, s), reverse=True)[:2]
214    ordered = [s for s in sentences if s in picked]  # keep the document's own order
215    return f"{' '.join(ordered)} [{best.id}]"
216
217
218def llm_answer(query: str, hits: list[Chunk]) -> str:
219    """Ask a model to answer from the retrieved chunks only, citing their ids."""
220    from openai import OpenAI  # optional dependency, only for --llm
221
222    model = os.environ.get("RAG_MODEL")
223    if not model:
224        sys.exit("Set RAG_MODEL to the model name your endpoint serves.")
225    sources = "\n\n".join(f"[{c.id}]\n{c.indexed}" for c in hits)
226    client = OpenAI(base_url=os.environ.get("OPENAI_BASE_URL"))
227    reply = client.chat.completions.create(
228        model=model,
229        temperature=0,
230        messages=[
231            {
232                "role": "system",
233                "content": "Answer only from the sources. Cite the id in brackets after each claim. "
234                "If the sources do not contain the answer, say you could not find it.",
235            },
236            {"role": "user", "content": f"Sources:\n{sources}\n\nQuestion: {query}"},
237        ],
238    )
239    return reply.choices[0].message.content or ""
240
241
242def main() -> None:
243    parser = argparse.ArgumentParser(description="Ask a question of the policy documents in ./corpus.")
244    parser.add_argument("question")
245    parser.add_argument("--k", type=int, default=3, help="chunks to retrieve (default 3)")
246    parser.add_argument("--mode", choices=["hybrid", "bm25", "ngram"], default="hybrid")
247    parser.add_argument("--chunking", choices=["headings", "windows"], default="headings")
248    parser.add_argument("--llm", action="store_true", help="answer with a model instead of quoting")
249    args = parser.parse_args()
250
251    chunks = load_chunks() if args.chunking == "headings" else load_windows()
252    hits = Retriever(chunks).search(args.question, k=args.k, mode=args.mode)
253    print(llm_answer(args.question, hits) if args.llm else extractive_answer(args.question, hits))
254    print("\nRetrieved:")
255    for position, chunk in enumerate(hits, 1):
256        print(f"  {position}. {chunk.id}")
257
258
259if __name__ == "__main__":
260    main()