"""mini_rag.py — the retrieval pieces built in lessons 06–10, collected so later lessons can import them. from site_docs import load_sections from mini_rag import HybridIndex index = HybridIndex(load_sections()) for hit in index.search("why is my spark job slow", k=3): print(hit.score, hit.doc["url"]) Nothing here needs a network or an API key: * BM25 keyword relevance, written from scratch (lesson 08) * LsaEmbedder dense vectors from TF-IDF + truncated SVD ("latent semantic analysis"). It stands in for a neural embedding model so the examples run offline. A real model (text-embedding-3-small, bge, e5…) captures meaning much better (lesson 06). * rrf reciprocal rank fusion, to merge several rankings (lesson 08) * GOLDEN questions with the lesson(s) that answer them, for evaluation (lesson 10) """ from __future__ import annotations import math import re from collections import Counter from dataclasses import dataclass import numpy as np TOKEN = re.compile(r"[a-z0-9]+") STOP = frozenset("a an the and or of to in on for is are was be by with as at it this that from how do i " "my what when why does doesn t s can you your we".split()) def tokenize(text: str) -> list[str]: return [t for t in TOKEN.findall(text.lower()) if t not in STOP] class BM25: """Okapi BM25 over a list of texts. scores(query) returns one score per text.""" def __init__(self, texts: list[str], k1: float = 1.5, b: float = 0.75) -> None: self.k1, self.b = k1, b docs = [Counter(tokenize(t)) for t in texts] self.lengths = np.array([sum(d.values()) for d in docs], dtype=float) self.avg_len = self.lengths.mean() df = Counter(term for d in docs for term in d) n = len(docs) self.idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()} self.postings: dict[str, list[tuple[int, int]]] = {} # term -> [(doc, term frequency)] for i, d in enumerate(docs): for term, tf in d.items(): self.postings.setdefault(term, []).append((i, tf)) def scores(self, query: str) -> np.ndarray: s = np.zeros(len(self.lengths)) for term in set(tokenize(query)): for i, tf in self.postings.get(term, ()): norm = 1 - self.b + self.b * self.lengths[i] / self.avg_len s[i] += self.idf[term] * tf * (self.k1 + 1) / (tf + self.k1 * norm) return s class LsaEmbedder: """Dense vectors without a neural model: TF-IDF, then SVD down to `dims` directions.""" def __init__(self, texts: list[str], dims: int = 256, seed: int = 0) -> None: from sklearn.decomposition import TruncatedSVD from sklearn.feature_extraction.text import TfidfVectorizer self.tfidf = TfidfVectorizer(tokenizer=tokenize, lowercase=False, token_pattern=None, sublinear_tf=True, min_df=2) self.svd = TruncatedSVD(n_components=dims, random_state=seed) self.matrix = self._unit(self.svd.fit_transform(self.tfidf.fit_transform(texts))) @staticmethod def _unit(m: np.ndarray) -> np.ndarray: return m / np.maximum(np.linalg.norm(m, axis=1, keepdims=True), 1e-12) def embed(self, texts: list[str]) -> np.ndarray: return self._unit(self.svd.transform(self.tfidf.transform(texts))) def scores(self, query: str) -> np.ndarray: return self.matrix @ self.embed([query])[0] # cosine: all vectors are unit length def rrf(*rankings: list[int], k: int = 60) -> list[int]: """Reciprocal rank fusion: score(d) = sum over rankings of 1 / (k + rank). Needs no score tuning.""" fused: Counter = Counter() for ranking in rankings: for rank, i in enumerate(ranking, start=1): fused[i] += 1 / (k + rank) return [i for i, _ in fused.most_common()] @dataclass class Hit: doc: dict score: float rank: int def doc_text(d: dict) -> str: """What gets indexed: the title and section heading carry a lot of signal, so include them.""" return f"{d['title']} — {d['section']}\n{d['text']}" class HybridIndex: def __init__(self, docs: list[dict], dims: int = 256) -> None: self.docs = docs texts = [doc_text(d) for d in docs] self.bm25 = BM25(texts) self.dense = LsaEmbedder(texts, dims=dims) def rank(self, query: str, mode: str = "hybrid", depth: int = 50) -> list[int]: bm = list(np.argsort(-self.bm25.scores(query))[:depth]) if mode == "bm25": return bm de = list(np.argsort(-self.dense.scores(query))[:depth]) return de if mode == "dense" else rrf(bm, de) def search(self, query: str, k: int = 5, mode: str = "hybrid", where: dict | None = None) -> list[Hit]: order = self.rank(query, mode, depth=200 if where else 50) if where: # metadata filter, e.g. {"course": "sql"} order = [i for i in order if all(self.docs[i].get(f) == v for f, v in where.items())] return [Hit(self.docs[i], round(1 / (r + 60), 5), r) for r, i in enumerate(order[:k], start=1)] def lessons_in(hits_or_indices, docs: list[dict] | None = None, k: int = 5) -> list[str]: """Collapse section hits to distinct lesson files, keeping rank order.""" seen: list[str] = [] for h in hits_or_indices: d = h.doc if isinstance(h, Hit) else docs[h] lesson = d["url"].split("#")[0] if lesson not in seen: seen.append(lesson) if len(seen) == k: break return seen # question -> lessons that answer it. Half share words with the lesson; half are paraphrased. GOLDEN: list[tuple[str, set[str]]] = [ ("Why doesn't my Spark code do anything until I call count()?", {"course/06-execution-model.html"}), ("How do I read the Spark UI to find a slow stage?", {"course/22-spark-ui.html"}), ("One key has most of the rows and a single task runs forever", {"course/23-skew-aqe.html"}), ("When should I cache a DataFrame?", {"course/20-caching.html"}), ("pandas UDF vs normal python UDF speed", {"course/16-udfs.html"}), ("Why does comparing a column to NULL return no rows?", {"sql/05-filtering-nulls.html"}), ("running total per customer ordered by date", {"sql/14-windows-1.html", "sql/15-windows-2.html"}), ("find customers who never placed an order", {"sql/10-advanced-joins.html", "sql/12-subqueries.html"}), ("walk an employee-manager hierarchy with a recursive query", {"sql/13-cte.html"}), ("fast pagination without OFFSET on large tables", {"sql/06-order-limit.html"}), ("dirty reads, phantom reads and isolation levels", {"sql/25-transactions.html", "mysql/08-mvcc.html"}), ("how does a B-tree index make lookups fast", {"sql/26-indexes.html", "mysql/10-btree.html"}), ("two transactions wait for each other forever in InnoDB", {"mysql/09-locking.html"}), ("composite index column order leftmost prefix rule", {"mysql/11-composite-indexes.html"}), ("set up a replica for high availability", {"mysql/19-replication.html"}), ("restore the database to a point in time using binary logs", {"mysql/18-backup.html"}), ("what is a closure and the nonlocal keyword", {"python/02-scope-closures.html"}), ("write a decorator that takes arguments", {"python/05-decorators.html"}), ("process a huge file lazily without loading it into memory", {"python/06-iterators-generators.html"}), ("why don't threads speed up CPU bound Python code", {"python/16-gil-concurrency.html"}), ("run many HTTP requests concurrently with the event loop", {"python/17-asyncio.html"}), ("find the shortest route in a weighted graph", {"dsa/18-shortest-paths.html", "dsa/25-project.html"}), ("autocomplete words by prefix", {"dsa/16-tries.html"}), ("keep the k largest items efficiently", {"dsa/15-heaps.html"}), ("govern tables, permissions and lineage across workspaces", {"databricks/06-unity-catalog.html"}), ("incrementally ingest new files landing in cloud storage", {"databricks/08-ingestion.html"}), ("bronze silver gold layers", {"databricks/09-medallion.html", "course/26-lakehouse.html"}), ("make the model return JSON that matches a Pydantic schema", {"openai-sdk/06-structured-outputs.html"}), ("the model asks my code to run a function, how do I loop", {"openai-sdk/07-function-calling.html"}), ("handle rate limit 429 errors with backoff", {"openai-sdk/12-errors-retries.html"}), ("show tokens to the user as they are generated", {"openai-sdk/05-streaming.html"}), ("reduce cost with prompt caching and batching", {"openai-sdk/13-cost-latency.html"}), ]