feat(ai): BM25 course retrieval over lessons, transcripts, code, deck text and notes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

packages/feature-ai/src/index.ts

@@ -4,3 +4,5 @@ export { createSecretBox, type SecretBox } from "./crypto";
 export const migrationsDirs: string[] = [];
 export * from "./settings";
 export * from "./providers";
+export * from "./retrieval/course-index";
+export { Bm25, tokenize } from "./retrieval/bm25";

packages/feature-ai/src/retrieval/bm25.ts

@@ -0,0 +1,79 @@
+const STOP = new Set(
+  "a an and are as at be but by for from has have i if in into is it its of on or so that the their then there these this to was were what when which who why will with you your".split(" "),
+);
+
+/** Light plural folding so "apples" matches "apple". */
+const stem = (w: string): string =>
+  w.length > 4 && w.endsWith("ies") ? `${w.slice(0, -3)}y` : w.length > 3 && w.endsWith("s") && !w.endsWith("ss") ? w.slice(0, -1) : w;
+
+/** Lowercased word tokens without stopwords; camelCase and snake_case identifiers also add their parts. */
+export function tokenize(text: string): string[] {
+  const out: string[] = [];
+  for (const raw of text.match(/[\p{L}\p{N}_]+/gu) ?? []) {
+    const lower = raw.toLowerCase();
+    if (STOP.has(lower) || lower.length > 64) continue;
+    out.push(stem(lower));
+    const parts = raw.split(/_|(?<=[\p{Ll}\p{N}])(?=\p{Lu})/u).filter(Boolean);
+    if (parts.length > 1) for (const p of parts) if (p.length > 1) out.push(stem(p.toLowerCase()));
+  }
+  return out;
+}
+
+type Doc<T> = { item: T; tf: Map<string, number>; len: number };
+
+const termFreq = (tokens: string[]) => {
+  const tf = new Map<string, number>();
+  for (const t of tokens) tf.set(t, (tf.get(t) ?? 0) + 1);
+  return tf;
+};
+
+/** Okapi BM25 (k1 = 1.2, b = 0.75) over an in-memory document list. */
+export class Bm25<T> {
+  private docs: Doc<T>[] = [];
+  private df = new Map<string, number>();
+  private totalLen = 0;
+
+  add(text: string, item: T): void {
+    const tokens = tokenize(text);
+    const tf = termFreq(tokens);
+    for (const t of tf.keys()) this.df.set(t, (this.df.get(t) ?? 0) + 1);
+    this.docs.push({ item, tf, len: tokens.length });
+    this.totalLen += tokens.length;
+  }
+
+  get size(): number { return this.docs.length; }
+
+  private idf(term: string): number {
+    const n = this.docs.length;
+    const df = this.df.get(term) ?? 0;
+    return Math.log(1 + (n - df + 0.5) / (df + 0.5));
+  }
+
+  private score(q: string[], tf: Map<string, number>, len: number): number {
+    const avg = this.docs.length ? this.totalLen / this.docs.length : Math.max(1, len);
+    let s = 0;
+    for (const term of q) {
+      const f = tf.get(term);
+      if (!f) continue;
+      s += this.idf(term) * ((f * 2.2) / (f + 1.2 * (0.25 + (0.75 * len) / avg)));
+    }
+    return s;
+  }
+
+  /** Top-k matches. `extra` documents are scored with this index's statistics without being added to it. */
+  search(query: string, k: number, extra: { text: string; item: T }[] = []): { item: T; score: number }[] {
+    const q = [...new Set(tokenize(query))];
+    if (!q.length) return [];
+    const results: { item: T; score: number }[] = [];
+    for (const d of this.docs) {
+      const s = this.score(q, d.tf, d.len);
+      if (s > 0) results.push({ item: d.item, score: s });
+    }
+    for (const e of extra) {
+      const tokens = tokenize(e.text);
+      const s = this.score(q, termFreq(tokens), tokens.length);
+      if (s > 0) results.push({ item: e.item, score: s });
+    }
+    return results.sort((a, b) => b.score - a.score).slice(0, k);
+  }
+}

packages/feature-ai/src/retrieval/course-index.ts

@@ -0,0 +1,236 @@
+import { readdir, stat } from "node:fs/promises";
+import { basename, extname, join, relative, sep } from "node:path";
+import { slugify, type LoadedCourse } from "@stc/course-format";
+import { Bm25 } from "./bm25";
+
+export type ChunkKind = "lesson" | "transcript" | "markdown" | "code" | "deck" | "note";
+export interface Chunk {
+  path: string;            // course-relative path (for cache text: the deck or lesson it belongs to)
+  kind: ChunkKind;
+  text: string;
+  heading?: string;
+  lessonId?: string;
+  lessonTitle?: string;
+}
+export interface CourseIndex { size: number; search(query: string, k: number, extra?: Chunk[]): Chunk[] }
+
+const SKIP_DIRS = new Set([".git", "node_modules", "dist", "build", "target", ".astro", ".next", "coverage", "bin", "obj", ".venv", "__pycache__"]);
+const MD = new Set([".md", ".markdown"]);
+const CODE = new Set([".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx", ".py", ".rs", ".scala", ".sc", ".fs", ".fsx", ".go", ".java", ".kt", ".rb", ".php", ".c", ".h", ".cpp", ".cs", ".sh", ".bash", ".ps1", ".sql", ".yaml", ".yml", ".toml", ".json", ".tf", ".html", ".css"]);
+const CODE_NAMES = new Set(["Dockerfile", "Makefile", "docker-compose.yml", "Containerfile"]);
+const CACHE_TEXT = new Set([".txt", ".md", ".vtt", ".srt"]);
+const SKIP_FILES = new Set(["package-lock.json", "bun.lock", "yarn.lock", "pnpm-lock.yaml", "Cargo.lock"]);
+const MAX_FILE = 200_000;
+const MAX_FILES = 2_000;
+const MAX_CHUNK = 1_500;
+const SOLUTION_SEGMENTS = new Set(["solution", "solutions", "reference"]);
+
+const posix = (p: string) => p.split(sep).join("/");
+
+/**
+ * Files only authors may see: manifest `browser.solution` files and anything under a `solution`, `solutions`
+ * or `reference` folder (Pure Snake `js/solutions/`, Systemic Design `reference/` and `expansion/solution/`).
+ * The assistant never retrieves or reads them.
+ */
+export function isSolutionPath(course: LoadedCourse, rel: string): boolean {
+  const p = posix(rel).replace(/^\.\//, "");
+  if (p.split("/").slice(0, -1).some((s) => SOLUTION_SEGMENTS.has(s.toLowerCase()))) return true;
+  if (/(^|\/)solution\.[a-z]+$/i.test(p)) return true;
+  for (const m of course.manifest?.modules ?? []) for (const l of m.lessons) if (l.browser?.solution && posix(l.browser.solution) === p) return true;
+  return false;
+}
+
+async function walk(root: string, onFile: (abs: string, rel: string) => Promise<boolean>, depth = 0, base = root): Promise<boolean> {
+  if (depth > 8) return true;
+  let entries;
+  try { entries = await readdir(root, { withFileTypes: true }); } catch { return true; }
+  entries.sort((a, b) => a.name.localeCompare(b.name));
+  for (const e of entries) {
+    if (e.name.startsWith(".") && e.name !== ".") continue;
+    const abs = join(root, e.name);
+    if (e.isDirectory()) {
+      if (SKIP_DIRS.has(e.name)) continue;
+      if (!(await walk(abs, onFile, depth + 1, base))) return false;
+    } else if (e.isFile()) {
+      if (!(await onFile(abs, posix(relative(base, abs))))) return false;
+    }
+  }
+  return true;
+}
+
+/** Splits Markdown at headings (outside code fences), then long sections at paragraph breaks. */
+export function splitMarkdown(text: string): { heading?: string; text: string }[] {
+  const sections: { heading?: string; lines: string[] }[] = [{ lines: [] }];
+  let fence = false;
+  for (const line of text.split("\n")) {
+    if (/^\s*(```|~~~)/.test(line)) fence = !fence;
+    const h = !fence && /^(#{1,6})\s+(.*)$/.exec(line);
+    if (h) sections.push({ heading: h[2]!.replace(/#+\s*$/, "").trim(), lines: [line] });
+    else sections[sections.length - 1]!.lines.push(line);
+  }
+  const out: { heading?: string; text: string }[] = [];
+  for (const s of sections) {
+    const body = s.lines.join("\n").trim();
+    if (!body) continue;
+    out.push(...splitLong(body).map((t) => ({ heading: s.heading, text: t })));
+  }
+  return out;
+}
+
+function splitLong(text: string): string[] {
+  if (text.length <= MAX_CHUNK) return [text];
+  const out: string[] = [];
+  let cur = "";
+  for (const para of text.split(/\n{2,}/)) {
+    if (cur && cur.length + para.length > MAX_CHUNK) { out.push(cur); cur = ""; }
+    if (para.length > MAX_CHUNK) {
+      for (let i = 0; i < para.length; i += MAX_CHUNK) out.push(para.slice(i, i + MAX_CHUNK));
+    } else cur = cur ? `${cur}\n\n${para}` : para;
+  }
+  if (cur) out.push(cur);
+  return out;
+}
+
+type LessonRef = { lessonId: string; lessonTitle: string; anchor?: string; transcript: boolean };
+
+function lessonRefs(course: LoadedCourse): Map<string, LessonRef[]> {
+  const refs = new Map<string, LessonRef[]>();
+  const add = (p: string | undefined, r: LessonRef) => {
+    if (!p) return;
+    const key = posix(p).replace(/^\.\//, "");
+    refs.set(key, [...(refs.get(key) ?? []), r]);
+  };
+  for (const m of course.manifest?.modules ?? []) {
+    for (const l of m.lessons) {
+      const base = { lessonId: l.id, lessonTitle: l.title };
+      if (l.type !== "video") add(l.file, { ...base, anchor: l.anchor, transcript: false });
+      add(l.transcript, { ...base, transcript: true });
+      if (l.hints?.file && l.hints.file !== l.file) add(l.hints.file, { ...base, transcript: false });
+      for (const t of Object.values(l.tracks ?? {})) add(t.file, { ...base, transcript: false });
+    }
+  }
+  return refs;
+}
+
+function pickLesson(refs: LessonRef[] | undefined, heading?: string): LessonRef | undefined {
+  if (!refs?.length) return undefined;
+  if (heading) {
+    const slug = slugify(heading);
+    const hit = refs.find((r) => r.anchor && r.anchor === slug);
+    if (hit) return hit;
+  }
+  return refs.find((r) => !r.anchor) ?? refs[0];
+}
+
+function lessonTitleOf(course: LoadedCourse, lessonId: string | undefined): string | undefined {
+  return lessonId ? course.lessonIndex.get(lessonId)?.lesson.title : undefined;
+}
+
+export async function buildCourseIndex(course: LoadedCourse, dataDir: string): Promise<CourseIndex> {
+  const bm = new Bm25<Chunk>();
+  const seen = new Set<string>();
+  const refs = lessonRefs(course);
+  const codeDirs = (course.manifest?.modules ?? []).flatMap((m) => (m.code ? [posix(m.code.dir).replace(/\/$/, "") + "/"] : []));
+  const add = (c: Chunk) => {
+    const key = `${c.kind}\u0000${c.text}`;
+    if (seen.has(key)) return;
+    seen.add(key);
+    bm.add([c.heading ?? "", c.lessonTitle ?? "", c.path, c.text].join("\n"), c);
+  };
+
+  let files = 0;
+  if (course.status === "ok") {
+    await walk(course.root, async (abs, rel) => {
+      if (files >= MAX_FILES) return false;
+      const ext = extname(rel).toLowerCase();
+      const name = basename(rel);
+      const isMd = MD.has(ext);
+      const isCode = !isMd && (CODE.has(ext) || CODE_NAMES.has(name) || codeDirs.some((d) => rel.startsWith(d)));
+      if ((!isMd && !isCode) || SKIP_FILES.has(name) || isSolutionPath(course, rel)) return true;
+      let size = 0;
+      try { size = (await stat(abs)).size; } catch { return true; }
+      if (size > MAX_FILE || size === 0) return true;
+      const text = await Bun.file(abs).text();
+      if (isCode && text.includes("\u0000")) return true;
+      files++;
+      const r = refs.get(rel);
+      if (isMd) {
+        for (const part of splitMarkdown(text)) {
+          const ref = pickLesson(r, part.heading);
+          add({
+            path: rel,
+            kind: ref ? (ref.transcript ? "transcript" : "lesson") : "markdown",
+            text: part.text,
+            heading: part.heading,
+            lessonId: ref?.lessonId,
+            lessonTitle: ref?.lessonTitle,
+          });
+        }
+      } else {
+        const lines = text.split("\n");
+        for (let i = 0; i < lines.length; i += 60) {
+          add({ path: rel, kind: "code", text: lines.slice(i, i + 60).join("\n"), heading: lines.length > 60 ? `lines ${i + 1}-${Math.min(lines.length, i + 60)}` : undefined });
+        }
+      }
+      return true;
+    });
+  }
+
+  // Derived text from the importer and transcription jobs: <dataDir>/cache/<courseId>/**.
+  const decks = new Map((course.manifest?.decks ?? []).map((d) => [d.id, d]));
+  const cacheRoot = join(dataDir, "cache", course.id);
+  await walk(cacheRoot, async (abs, rel) => {
+    if (!CACHE_TEXT.has(extname(rel).toLowerCase())) return true;
+    const size = (await stat(abs).catch(() => null))?.size ?? 0;
+    if (!size || size > MAX_FILE * 5) return true;
+    const text = await Bun.file(abs).text();
+    const segs = rel.split("/");
+    const deck = segs.map((s) => decks.get(s)).find(Boolean);
+    const lessonSeg = segs.map((s) => s.replace(/\.[^.]+$/, "")).find((s) => course.lessonIndex.has(s));
+    const page = /(?:page|slide)[-_ ]?(\d+)/i.exec(basename(rel))?.[1];
+    if (deck) {
+      const lessonId = course.lessonIndex.has(deck.attachTo) ? deck.attachTo : undefined;
+      for (const part of splitLong(text.trim())) {
+        add({ path: posix(deck.file), kind: "deck", text: part, heading: page ? `Slide ${Number(page)}` : deck.title, lessonId, lessonTitle: lessonTitleOf(course, lessonId) });
+      }
+    } else {
+      const kind: ChunkKind = /transcript|\.vtt$|\.srt$/i.test(rel) || lessonSeg ? "transcript" : "deck";
+      for (const part of splitLong(text.trim())) {
+        add({ path: `cache/${rel}`, kind, text: part, lessonId: lessonSeg, lessonTitle: lessonTitleOf(course, lessonSeg) });
+      }
+    }
+    return true;
+  });
+
+  return {
+    get size() { return bm.size; },
+    search(query, k, extra = []) {
+      return bm.search(query, k, extra.map((c) => ({ text: [c.heading ?? "", c.lessonTitle ?? "", c.text].join("\n"), item: c }))).map((r) => r.item);
+    },
+  };
+}
+
+/** The user's notes as chunks, scored at query time and never cached. */
+export function noteChunks(course: LoadedCourse, notes: { id: string; lessonId: string; body: string; anchorKind: string; anchorValue: string | null }[]): Chunk[] {
+  return notes.map((n) => ({
+    path: "notes",
+    kind: "note" as const,
+    text: n.body,
+    heading: n.anchorKind === "video" && n.anchorValue ? `at ${n.anchorValue}s` : n.anchorKind === "slide" && n.anchorValue ? `slide ${n.anchorValue}` : undefined,
+    lessonId: n.lessonId,
+    lessonTitle: lessonTitleOf(course, n.lessonId),
+  }));
+}
+
+const TTL_MS = 5 * 60_000;
+const cache = new WeakMap<LoadedCourse, { at: number; index: Promise<CourseIndex> }>();
+
+/** Cached per LoadedCourse object: a repository reload yields new objects, and entries expire after 5 minutes. */
+export function getCourseIndex(course: LoadedCourse, dataDir: string): Promise<CourseIndex> {
+  const hit = cache.get(course);
+  if (hit && Date.now() - hit.at < TTL_MS) return hit.index;
+  const index = buildCourseIndex(course, dataDir);
+  cache.set(course, { at: Date.now(), index });
+  index.catch(() => cache.delete(course));
+  return index;
+}

packages/feature-ai/test/retrieval.test.ts

@@ -0,0 +1,78 @@
+import { afterAll, beforeAll, expect, test } from "bun:test";
+import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { createCourseRepository, type LoadedCourse } from "@stc/course-format";
+import { Bm25, tokenize } from "../src/retrieval/bm25";
+import { buildCourseIndex, getCourseIndex, isSolutionPath, noteChunks } from "../src/retrieval/course-index";
+
+const FIXTURES = join(import.meta.dir, "../../../fixtures/courses");
+let course: LoadedCourse;
+let dataDir: string;
+
+beforeAll(async () => {
+  const repo = await createCourseRepository([FIXTURES]);
+  course = repo.get("sample")!;
+  dataDir = mkdtempSync(join(tmpdir(), "stc-ai-"));
+  mkdirSync(join(dataDir, "cache/sample/decks/d1"), { recursive: true });
+  writeFileSync(join(dataDir, "cache/sample/decks/d1/page-3.txt"), "Container images are layered filesystems built from a Dockerfile.");
+  mkdirSync(join(dataDir, "cache/sample/transcripts"), { recursive: true });
+  writeFileSync(join(dataDir, "cache/sample/transcripts/m1-video.txt"), "In this clip the presenter explains the test pattern calibration.");
+});
+afterAll(() => rmSync(dataDir, { recursive: true, force: true }));
+
+test("tokenize lowercases, drops stopwords and splits identifiers", () => {
+  expect(tokenize("The inBounds function, and dropLast!")).toEqual(["inbound", "in", "bound", "function", "droplast", "drop", "last"]);
+  expect(tokenize("Æble på dansk")).toEqual(["æble", "på", "dansk"]);
+});
+
+test("BM25 ranks the more relevant document first", () => {
+  const b = new Bm25<string>();
+  b.add("snake moves one cell per tick", "a");
+  b.add("apples spawn at random free cells", "b");
+  b.add("the snake grows when it eats an apple apple", "c");
+  expect(b.search("apple", 5).map((r) => r.item)).toEqual(["c", "b"]);
+  expect(b.search("zebra", 5)).toEqual([]);
+});
+
+test("finds lesson text and maps it to a lesson", async () => {
+  const idx = await buildCourseIndex(course, dataDir);
+  const hits = idx.search("sum of its arguments", 3);
+  expect(hits[0]!.path).toBe("lessons/04-exercise.md");
+  expect(["m2-text", "m2-ex-add"]).toContain(hits[0]!.lessonId!);
+  expect(hits[0]!.lessonTitle).toBeTruthy();
+});
+
+test("indexes code, transcripts and deck text from the cache", async () => {
+  const idx = await buildCourseIndex(course, dataDir);
+  const code = idx.search("hello.sh", 5).find((c) => c.kind === "code");
+  expect(code?.path).toBe("code/S01 - First steps/hello.sh");
+  const deck = idx.search("Dockerfile layered images", 3)[0]!;
+  expect(deck.kind).toBe("deck");
+  expect(deck.path).toBe("media/deck.pdf");
+  expect(deck.heading).toBe("Slide 3");
+  expect(deck.lessonId).toBe("m1-video");
+  const tr = idx.search("calibration presenter", 3)[0]!;
+  expect(tr.kind).toBe("transcript");
+  expect(tr.lessonId).toBe("m1-video");
+});
+
+test("solution files are never indexed", async () => {
+  const idx = await buildCourseIndex(course, dataDir);
+  for (const q of ["add", "return a + b", "export const add"]) {
+    expect(idx.search(q, 50).some((c) => c.path.includes("solution"))).toBe(false);
+  }
+  expect(isSolutionPath(course, "exercises/add/solution.js")).toBe(true);
+  expect(isSolutionPath(course, "js/solutions/rules.js")).toBe(true);
+  expect(isSolutionPath(course, "course/01-handles/reference/ts/src/engine.ts")).toBe(true);
+  expect(isSolutionPath(course, "course/01-handles/expansion/solution/x.rs")).toBe(true);
+  expect(isSolutionPath(course, "exercises/add/starter.js")).toBe(false);
+});
+
+test("notes are searched alongside the course", async () => {
+  const idx = await getCourseIndex(course, dataDir);
+  expect(await getCourseIndex(course, dataDir)).toBe(idx);
+  const extra = noteChunks(course, [{ id: "n1", lessonId: "m1-intro", body: "Remember: zebras are striped horses", anchorKind: "none", anchorValue: null }]);
+  const hits = idx.search("zebras striped", 3, extra);
+  expect(hits[0]).toMatchObject({ kind: "note", lessonId: "m1-intro", lessonTitle: "Introduction" });
+});